<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

import (
	"bytes"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"
	"testing"

	"github.com/ghodss/yaml"

	"github.com/kubernetes-sigs/kustomize/pkg/fs"
	"k8s.io/apimachinery/pkg/util/sets"
)

type DiffTestCase struct {
	Description string   `yaml:"description"`
	Args        []string `yaml:"args"`
	Filename    string   `yaml:"filename"`
	// path to the file that contains the expected output
	ExpectedDiff  string `yaml:"expectedDiff"`
	ExpectedError string `yaml:"expectedError"`
}

func TestDiff(t *testing.T) {
	const updateEnvVar = "UPDATE_KUSTOMIZE_EXPECTED_DATA"
	updateKustomizeExpected := os.Getenv(updateEnvVar) == "true"

	noopDir, _ := regexp.Compile(`/tmp/noop-[0-9]*/`)
	transformedDir, _ := regexp.Compile(`/tmp/transformed-[0-9]*/`)
	timestamp, _ := regexp.Compile(`[0-9]{4}-(0[1-9]|1[0-2])-(0[1-9]|[1-2][0-9]|3[0-1]) (2[0-3]|[01][0-9]):[0-5][0-9]:[0-5][0-9].[0-9]* [+-]{1}[0-9]{4}`)

	fs := fs.MakeRealFS()

	testcases := sets.NewString()
	filepath.Walk("testdata", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if path == "testdata" {
			return nil
		}
		name := filepath.Base(path)
		if info.IsDir() {
			if strings.HasPrefix(name, "testcase-") {
				testcases.Insert(strings.TrimPrefix(name, "testcase-"))
			}
			return filepath.SkipDir
		}
		return nil
	})
	// sanity check that we found the right folder
	if !testcases.Has("simple") {
		t.Fatalf("Error locating testcases")
	}

	for _, testcaseName := range testcases.List() {
		t.Run(testcaseName, func(t *testing.T) {
			runDiffTestCase(t, testcaseName, updateKustomizeExpected, fs,
				noopDir, transformedDir, timestamp)
		})
	}
}

func runDiffTestCase(t *testing.T, testcaseName string, updateKustomizeExpected bool, fs fs.FileSystem,
	noopDir, transformedDir, timestamp *regexp.Regexp) {
	name := testcaseName
	testcase := DiffTestCase{}
	testcaseDir := filepath.Join("testdata", "testcase-"+name)
	testcaseData, err := ioutil.ReadFile(filepath.Join(testcaseDir, "test.yaml"))
	if err != nil {
		t.Fatalf("%s: %v", name, err)
	}
	if err := yaml.Unmarshal(testcaseData, &testcase); err != nil {
		t.Fatalf("%s: %v", name, err)
	}

	diffOps := &diffOptions{
		kustomizationPath: testcase.Filename,
	}
	buf := bytes.NewBuffer([]byte{})
	err = diffOps.RunDiff(buf, os.Stderr, fs)
	switch {
	case err != nil && len(testcase.ExpectedError) == 0:
		t.Errorf("unexpected error: %v", err)
	case err != nil && len(testcase.ExpectedError) != 0:
		if !strings.Contains(err.Error(), testcase.ExpectedError) {
			t.Errorf("expected error to contain %q but got: %v", testcase.ExpectedError, err)
		}
		return
	case err == nil && len(testcase.ExpectedError) != 0:
		t.Errorf("unexpected no error")
	}

	actualString := string(buf.Bytes())
	actualString = noopDir.ReplaceAllString(actualString, "/tmp/noop/")
	actualString = transformedDir.ReplaceAllString(actualString, "/tmp/transformed/")
	actualString = timestamp.ReplaceAllString(actualString, "YYYY-MM-DD HH:MM:SS")
	actualBytes := []byte(actualString)
	if !updateKustomizeExpected {
		expectedBytes, err := ioutil.ReadFile(testcase.ExpectedDiff)
		if err != nil {
			t.Errorf("unexpected error: %v", err)
		}
		if !reflect.DeepEqual(actualBytes, expectedBytes) {
			t.Errorf("%s\ndoesn't equal expected:\n%s\n", actualBytes, expectedBytes)
		}
	} else {
		ioutil.WriteFile(testcase.ExpectedDiff, actualBytes, 0644)
	}

}
<commit_msg>Fix wrong path in diff tests on macOS<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

import (
	"bytes"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"
	"testing"

	"github.com/ghodss/yaml"

	"github.com/kubernetes-sigs/kustomize/pkg/fs"
	"k8s.io/apimachinery/pkg/util/sets"
)

type DiffTestCase struct {
	Description string   `yaml:"description"`
	Args        []string `yaml:"args"`
	Filename    string   `yaml:"filename"`
	// path to the file that contains the expected output
	ExpectedDiff  string `yaml:"expectedDiff"`
	ExpectedError string `yaml:"expectedError"`
}

func TestDiff(t *testing.T) {
	const updateEnvVar = "UPDATE_KUSTOMIZE_EXPECTED_DATA"
	updateKustomizeExpected := os.Getenv(updateEnvVar) == "true"

	tempDir := regexp.QuoteMeta(filepath.Clean(os.TempDir()))
	noopDir, _ := regexp.Compile(tempDir + `/noop-[0-9]*/`)
	transformedDir, _ := regexp.Compile(tempDir + `/transformed-[0-9]*/`)
	timestamp, _ := regexp.Compile(`[0-9]{4}-(0[1-9]|1[0-2])-(0[1-9]|[1-2][0-9]|3[0-1]) (2[0-3]|[01][0-9]):[0-5][0-9]:[0-5][0-9].[0-9]* [+-]{1}[0-9]{4}`)

	fs := fs.MakeRealFS()

	testcases := sets.NewString()
	filepath.Walk("testdata", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if path == "testdata" {
			return nil
		}
		name := filepath.Base(path)
		if info.IsDir() {
			if strings.HasPrefix(name, "testcase-") {
				testcases.Insert(strings.TrimPrefix(name, "testcase-"))
			}
			return filepath.SkipDir
		}
		return nil
	})
	// sanity check that we found the right folder
	if !testcases.Has("simple") {
		t.Fatalf("Error locating testcases")
	}

	for _, testcaseName := range testcases.List() {
		t.Run(testcaseName, func(t *testing.T) {
			runDiffTestCase(t, testcaseName, updateKustomizeExpected, fs,
				noopDir, transformedDir, timestamp)
		})
	}
}

func runDiffTestCase(t *testing.T, testcaseName string, updateKustomizeExpected bool, fs fs.FileSystem,
	noopDir, transformedDir, timestamp *regexp.Regexp) {
	name := testcaseName
	testcase := DiffTestCase{}
	testcaseDir := filepath.Join("testdata", "testcase-"+name)
	testcaseData, err := ioutil.ReadFile(filepath.Join(testcaseDir, "test.yaml"))
	if err != nil {
		t.Fatalf("%s: %v", name, err)
	}
	if err := yaml.Unmarshal(testcaseData, &testcase); err != nil {
		t.Fatalf("%s: %v", name, err)
	}

	diffOps := &diffOptions{
		kustomizationPath: testcase.Filename,
	}
	buf := bytes.NewBuffer([]byte{})
	err = diffOps.RunDiff(buf, os.Stderr, fs)
	switch {
	case err != nil && len(testcase.ExpectedError) == 0:
		t.Errorf("unexpected error: %v", err)
	case err != nil && len(testcase.ExpectedError) != 0:
		if !strings.Contains(err.Error(), testcase.ExpectedError) {
			t.Errorf("expected error to contain %q but got: %v", testcase.ExpectedError, err)
		}
		return
	case err == nil && len(testcase.ExpectedError) != 0:
		t.Errorf("unexpected no error")
	}

	actualString := string(buf.Bytes())
	actualString = noopDir.ReplaceAllString(actualString, "/tmp/noop/")
	actualString = transformedDir.ReplaceAllString(actualString, "/tmp/transformed/")
	actualString = timestamp.ReplaceAllString(actualString, "YYYY-MM-DD HH:MM:SS")
	actualBytes := []byte(actualString)
	if !updateKustomizeExpected {
		expectedBytes, err := ioutil.ReadFile(testcase.ExpectedDiff)
		if err != nil {
			t.Errorf("unexpected error: %v", err)
		}
		if !reflect.DeepEqual(actualBytes, expectedBytes) {
			t.Errorf("%s\ndoesn't equal expected:\n%s\n", actualBytes, expectedBytes)
		}
	} else {
		ioutil.WriteFile(testcase.ExpectedDiff, actualBytes, 0644)
	}

}
<|endoftext|>


<commit_before>package version

// Version is the version of the build.
const Version = "0.1.18-dev"
<commit_msg>bump to v0.1.18<commit_after>package version

// Version is the version of the build.
const Version = "0.1.18"
<|endoftext|>


<commit_before>package version

// VERSION ...
const VERSION = "1.1.58"
<commit_msg>v1.1.59 [skip ci]<commit_after>package version

// VERSION ...
const VERSION = "1.1.59"
<|endoftext|>


<commit_before>package version

import "fmt"

const (
	// VersionMajor is for an API incompatible changes
	VersionMajor = 0
	// VersionMinor is for functionality in a backwards-compatible manner
	VersionMinor = 1
	// VersionPatch is for backwards-compatible bug fixes
	VersionPatch = 0

	// VersionDev indicates development branch. Releases will be empty string.
	VersionDev = "-dev"
)

// Version is the specification version that the package types support.
var Version = fmt.Sprintf("%d.%d.%d%s", VersionMajor, VersionMinor, VersionPatch, VersionDev)
<commit_msg>Bump to version 1.4<commit_after>package version

import "fmt"

const (
	// VersionMajor is for an API incompatible changes
	VersionMajor = 0
	// VersionMinor is for functionality in a backwards-compatible manner
	VersionMinor = 1
	// VersionPatch is for backwards-compatible bug fixes
	VersionPatch = 4

	// VersionDev indicates development branch. Releases will be empty string.
	VersionDev = ""
)

// Version is the specification version that the package types support.
var Version = fmt.Sprintf("%d.%d.%d%s", VersionMajor, VersionMinor, VersionPatch, VersionDev)
<|endoftext|>


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

package version

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

package version

const (
	Version = "0.3.1+git"
)
<|endoftext|>


<commit_before>// The version package provides a location to set the release versions for all
// packages to consume, without creating import cycles.
//
// This pckage should not import any other terraform packages.
package version

import (
	"fmt"

	version "github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.11.8"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
var Prerelease = "dev"

// SemVer is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVer = version.Must(version.NewVersion(Version))

// Header is the header name used to send the current terraform version
// in http requests.
const Header = "Terraform-Version"

// String returns the complete version string, including prerelease
func String() string {
	if Prerelease != "" {
		return fmt.Sprintf("%s-%s", Version, Prerelease)
	}
	return Version
}
<commit_msg>Fix typo in version docs<commit_after>// The version package provides a location to set the release versions for all
// packages to consume, without creating import cycles.
//
// This package should not import any other terraform packages.
package version

import (
	"fmt"

	version "github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.11.8"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
var Prerelease = "dev"

// SemVer is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVer = version.Must(version.NewVersion(Version))

// Header is the header name used to send the current terraform version
// in http requests.
const Header = "Terraform-Version"

// String returns the complete version string, including prerelease
func String() string {
	if Prerelease != "" {
		return fmt.Sprintf("%s-%s", Version, Prerelease)
	}
	return Version
}
<|endoftext|>


<commit_before>/*
NotifyDispatcher attempts to make working with a single Listener easy for a
dynamic set of independent listeners.

Usage

Viz:

    import (
        "github.com/lib/pq"
        "github.com/johto/notifyutils/notifydispatcher"
        "fmt"
        "time"
    )

    func listener(dispatcher *notifydispatcher.NotifyDispatcher) {
        ch := make(chan *pq.Notification, 8)
        err := dispatcher.Listen("listenerchannel", ch)
        if err != nil {
            panic(err)
        }
        for n := range ch {
            if n == nil {
                fmt.Println("lost connection, but we're fine now!")
                continue
            }

            fmt.Println("received notification!")
            // do something with notification
        }
        panic("could not keep up!")
    }

    func main() {
        dispatcher := notifydispatcher.NewNotifyDispatcher(pq.NewListener("", time.Second, time.Minute, nil))
        for i := 0; i < 8; i++ {
            go listener(dispatcher)
        }
        select{}
    }

*/
package notifydispatcher

import (
	"container/list"
	"errors"
	"fmt"
	"github.com/lib/pq"
	"sync"
)

var (
	ErrChannelAlreadyActive = errors.New("channel is already active")
	ErrChannelNotActive = errors.New("channel is not active")
)

var errClosed = errors.New("NotifyDispatcher has been closed")

// SlowReaderEliminationStrategy controls the behaviour of the dispatcher in
// case the buffer of a listener's channel is full and attempting to send to it
// would block the dispatcher, preventing it from delivering notifications for
// unrelated listeners.  The default is CloseSlowReaders, but it can be changed
// at any point during a dispatcher's lifespan using
// SetSlowReaderEliminationStrategy.
type SlowReaderEliminationStrategy int
const (
	// When a send would block, the listener's channel is removed from the set
	// of listeners for that notification channel, and the channel is closed.
	// This is the default strategy.
	CloseSlowReaders SlowReaderEliminationStrategy = iota

	// When a send would block, the notification is not delivered.  Delivery is
	// not reattempted.
	NeglectSlowReaders
)

type listenRequest struct {
	channel string
	unlisten bool
}

type BroadcastChannel struct {
	Channel chan struct{}
	elem *list.Element
}

type NotifyDispatcher struct {
	listener *pq.Listener

	// Some details about the behaviour.  Only touch or look at while holding
	// "lock".
	slowReaderEliminationStrategy SlowReaderEliminationStrategy
	broadcastOnConnectionLoss bool
	broadcastChannels *list.List

	listenRequestch chan listenRequest

	lock sync.Mutex
	channels map[string] *listenSet
	closed bool
	// provide an escape hatch for goroutines sending on listenRequestch
	closeChannel chan bool
}

// NewNotifyDispatcher creates a new NotifyDispatcher, using the supplied
// pq.Listener underneath.  The ownership of the Listener is transferred to
// NotifyDispatcher.  You should not use it after calling NewNotifyDispatcher.
func NewNotifyDispatcher(l *pq.Listener) *NotifyDispatcher {
	d := &NotifyDispatcher{
		listener: l,
		slowReaderEliminationStrategy: CloseSlowReaders,
		broadcastOnConnectionLoss: true,
		broadcastChannels: list.New(),
		listenRequestch: make(chan listenRequest, 64),
		channels: make(map[string] *listenSet),
		closeChannel: make(chan bool),
	}
	go d.dispatcherLoop()
	go d.listenRequestHandlerLoop()
	return d
}

// Sets the strategy for mitigating the adverse effects slow readers might have
// on the dispatcher.  See SlowReaderEliminationStrategy.
func (d *NotifyDispatcher) SetSlowReaderEliminationStrategy(strategy SlowReaderEliminationStrategy) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.slowReaderEliminationStrategy = strategy
}

// Controls whether a nil notification from the underlying Listener is
// broadcast to all channels in the set.
func (d *NotifyDispatcher) SetBroadcastOnConnectionLoss(value bool) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.broadcastOnConnectionLoss = value
}

// Opens a new "broadcast channel".  A broadcast channel is sent to by the
// NotifyDispatcher every time the underlying Listener has re-established its
// database connection.
func (d *NotifyDispatcher) OpenBroadcastChannel() BroadcastChannel {
	d.lock.Lock()
	defer d.lock.Unlock()

	ch := BroadcastChannel{
		Channel: make(chan struct{}, 1),
	}
	ch.elem = d.broadcastChannels.PushFront(ch)
	return ch
}

// Closes the broadcast channel ch.
func (d *NotifyDispatcher) CloseBroadcastChannel(ch BroadcastChannel) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.broadcastChannels.Remove(ch.elem) != ch {
		panic("oops")
	}
	close(ch.Channel)
}

// Broadcast on all broadcastChannels and on all channels unless
// broadcastOnConnectionLoss is disabled.
func (d *NotifyDispatcher) broadcast() {
	reapchans := []string{}

	d.lock.Lock()
	for e := d.broadcastChannels.Front(); e != nil; e = e.Next() {
		select {
			case e.Value.(BroadcastChannel).Channel <- struct{}{}:
			default:
				// nothing to do
		}
	}

	if !d.broadcastOnConnectionLoss {
		d.lock.Unlock()
		return
	}

	for channel, set := range d.channels {
		if !set.broadcast(d.slowReaderEliminationStrategy, nil) {
			reapchans = append(reapchans, channel)
		}
	}
	d.lock.Unlock()

	for _, ch := range reapchans {
		d.listenRequestch <- listenRequest{ch, true}
	}
}

func (d *NotifyDispatcher) dispatch(n *pq.Notification) {
	reap := false

	d.lock.Lock()
	set, ok := d.channels[n.Channel]
	if ok {
		reap = !set.broadcast(d.slowReaderEliminationStrategy, n)
	}
	d.lock.Unlock()

	if reap {
		d.listenRequestch <- listenRequest{n.Channel, true}
	}
}

func (d *NotifyDispatcher) dispatcherLoop() {
	for {
		n := <-d.listener.Notify
		if n == nil {
			d.broadcast()
		} else {
			d.dispatch(n)
		}
	}
}

// Attempt to start listening on channel.  The caller should not be holding
// lock.
func (d *NotifyDispatcher) execListen(channel string) {
	for {
		err := d.listener.Listen(channel)
		// ErrChannelAlreadyOpen is a valid return value here; we could have
		// abandoned a channel in Unlisten() if the server returned an error
		// for no apparent reason.
		if err == nil ||
		   err == pq.ErrChannelAlreadyOpen {
			break
		}
	}

	d.lock.Lock()
	defer d.lock.Unlock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	set.setState(listenSetStateActive)
}

func (d *NotifyDispatcher) execUnlisten(channel string) {
	// we don't really care about the error
	_ = d.listener.Unlisten(channel)

	d.lock.Lock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	if set.state != listenSetStateZombie {
		panic("oops")
	}
	if set.reap() {
		delete(d.channels, channel)
		d.lock.Unlock()
	} else {
		// Couldn't reap the set because it got new listeners while we were
		// waiting for the UNLISTEN to go through.  Re-LISTEN it, but remember
		// to release the lock first.
		d.lock.Unlock()

		d.execListen(channel)
	}
}

func (d *NotifyDispatcher) listenRequestHandlerLoop() {
	for {
		// check closeChannel, just in case we've been closed and there's a
		// backlog of requests
		select {
			case <-d.closeChannel:
				return
			default:
		}

		select {
			case <-d.closeChannel:
				return

			case req := <-d.listenRequestch:
				if req.unlisten {
					d.execUnlisten(req.channel)
				} else {
					d.execListen(req.channel)
				}
		}
	}
}

func (d *NotifyDispatcher) requestListen(channel string, unlisten bool) error {
	select {
		// make sure we don't get stuck here if someone Close()s us
		case <-d.closeChannel:
			return errClosed
		case d.listenRequestch <- listenRequest{channel, unlisten}:
			return nil
	}

	panic("not reached")
}

// Listen adds ch to the set of listeners for notification channel channel.  ch
// should be a buffered channel.  If ch is already in the set of listeners for
// channel, ErrChannelAlreadyActive is returned.  After Listen has returned,
// the notification channel is open and the dispatcher will attempt to deliver
// all notifications received for that channel to ch.
//
// If SlowReaderEliminationStrategy is CloseSlowReaders, reusing the same ch
// for multiple notification channels is not allowed.
func (d *NotifyDispatcher) Listen(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	set, ok := d.channels[channel]
	if ok {
		err := set.add(ch)
		if err != nil {
			return err
		}
	} else {
		set = d.newListenSet(ch)
		d.channels[channel] = set

		// must not be holding the lock while requesting a listen
		d.lock.Unlock()
		err := d.requestListen(channel, false)
		if err != nil {
			return err
		}
		d.lock.Lock()
	}

	return set.waitForActive(&d.closed)
}

// Removes ch from the set of listeners for notification channel channel.  If
// ch is not in the set of listeners for channel, ErrChannelNotActive is
// returned.
func (d *NotifyDispatcher) Unlisten(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()

	if d.closed {
		d.lock.Unlock()
		return errClosed
	}

	set, ok := d.channels[channel]
	if !ok {
		d.lock.Unlock()
		return ErrChannelNotActive
	}
	last, err := set.remove(ch)
	if err != nil {
		d.lock.Unlock()
		return err
	}
	if !last {
		// the set isn't empty; nothing further for us to do
		d.lock.Unlock()
		return nil
	}
	d.lock.Unlock()

	// we were the last listener, cue the reaper
	return d.requestListen(channel, true)
}

func (d *NotifyDispatcher) Close() error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	d.closed = true
	close(d.closeChannel)

	for _, set := range d.channels {
		set.activeOrClosedCond.Broadcast()
	}

	// TODO: kill the dispatcher as well

	return nil
}

type listenSetState int
const (
	// The set was recently spawned or respawned, and it's waiting for a call
	// to Listen() to succeed.
	listenSetStateNewborn listenSetState = iota
	// The set is ready and any notifications from the database will be
	// dispatched to the set.
	listenSetStateActive
	// The set has recently been emptied, and it's waiting for a call to
	// Unlisten() to finish.
	listenSetStateZombie
)

type listenSet struct {
	channels map[chan<- *pq.Notification] struct{}
	state listenSetState
	activeOrClosedCond *sync.Cond
}

func (d *NotifyDispatcher) newListenSet(firstInhabitant chan<- *pq.Notification) *listenSet {
	s := &listenSet{
		channels: make(map[chan<- *pq.Notification] struct{}),
		state: listenSetStateNewborn,
	}
	s.activeOrClosedCond = sync.NewCond(&d.lock)
	s.channels[firstInhabitant] = struct{}{}
	return s
}

func (s *listenSet) setState(newState listenSetState) {
	var expectedState listenSetState
	switch newState {
		case listenSetStateNewborn:
			expectedState = listenSetStateZombie
		case listenSetStateActive:
			expectedState = listenSetStateNewborn
		case listenSetStateZombie:
			expectedState = listenSetStateActive
	}
	if s.state != expectedState {
		panic(fmt.Sprintf("illegal state transition from %v to %v", s.state, newState))
	}
	s.state = newState
	if s.state == listenSetStateActive {
		s.activeOrClosedCond.Broadcast()
	}
}

func (s *listenSet) add(ch chan<- *pq.Notification) error {
	_, ok := s.channels[ch]
	if ok {
		return ErrChannelAlreadyActive
	}
	s.channels[ch] = struct{}{}
	return nil
}

func (s *listenSet) remove(ch chan<- *pq.Notification) (last bool, err error) {
	_, ok := s.channels[ch]
	if !ok {
		return false, ErrChannelNotActive
	}
	delete(s.channels, ch)

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return true, nil
	}
	return false, nil
}

func (s *listenSet) broadcast(strategy SlowReaderEliminationStrategy, n *pq.Notification) bool {
	// must be active
	if s.state != listenSetStateActive {
		return true
	}

	for ch := range s.channels {
		select {
			case ch <- n:

			default:
				if strategy == CloseSlowReaders {
					delete(s.channels, ch)
					close(ch)
				}
		}
	}

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return false
	}

	return true
}

// Marks the set active after a successful call to Listen().
func (s *listenSet) markActive() {
	s.setState(listenSetStateActive)
}

// Wait for the listen set to become "active".  Returns nil if successfull, or
// errClosed if the dispatcher was closed while waiting.  The caller should be
// holding d.lock.
func (s *listenSet) waitForActive(closed *bool) error {
	for {
		if *closed {
			return errClosed
		}
		if s.state == listenSetStateActive {
			return nil
		}
		s.activeOrClosedCond.Wait()
	}
}

// Try to reap a zombie set after Unlisten().  Returns true if the set should
// be removed, false otherwise.
func (s *listenSet) reap() bool {
	if s.state != listenSetStateZombie {
		panic("unexpected state in reap")
	}

	if len(s.channels) > 0 {
		// we need to be respawned
		s.setState(listenSetStateNewborn)
		return false
	}

	return true
}
<commit_msg>Add comments, improve code<commit_after>/*
NotifyDispatcher attempts to make working with a single Listener easy for a
dynamic set of independent listeners.

Usage

Viz:

    import (
        "github.com/lib/pq"
        "github.com/johto/notifyutils/notifydispatcher"
        "fmt"
        "time"
    )

    func listener(dispatcher *notifydispatcher.NotifyDispatcher) {
        ch := make(chan *pq.Notification, 8)
        err := dispatcher.Listen("listenerchannel", ch)
        if err != nil {
            panic(err)
        }
        for n := range ch {
            if n == nil {
                fmt.Println("lost connection, but we're fine now!")
                continue
            }

            fmt.Println("received notification!")
            // do something with notification
        }
        panic("could not keep up!")
    }

    func main() {
        dispatcher := notifydispatcher.NewNotifyDispatcher(pq.NewListener("", time.Second, time.Minute, nil))
        for i := 0; i < 8; i++ {
            go listener(dispatcher)
        }
        select{}
    }

*/
package notifydispatcher

import (
	"container/list"
	"errors"
	"fmt"
	"github.com/lib/pq"
	"sync"
)

var (
	ErrChannelAlreadyActive = errors.New("channel is already active")
	ErrChannelNotActive = errors.New("channel is not active")
)

var errClosed = errors.New("NotifyDispatcher has been closed")

// SlowReaderEliminationStrategy controls the behaviour of the dispatcher in
// case the buffer of a listener's channel is full and attempting to send to it
// would block the dispatcher, preventing it from delivering notifications for
// unrelated listeners.  The default is CloseSlowReaders, but it can be changed
// at any point during a dispatcher's lifespan using
// SetSlowReaderEliminationStrategy.
type SlowReaderEliminationStrategy int
const (
	// When a send would block, the listener's channel is removed from the set
	// of listeners for that notification channel, and the channel is closed.
	// This is the default strategy.
	CloseSlowReaders SlowReaderEliminationStrategy = iota

	// When a send would block, the notification is not delivered.  Delivery is
	// not reattempted.
	NeglectSlowReaders
)

type listenRequest struct {
	channel string
	unlisten bool
}

type BroadcastChannel struct {
	Channel chan struct{}
	elem *list.Element
}

type NotifyDispatcher struct {
	listener *pq.Listener

	// Some details about the behaviour.  Only touch or look at while holding
	// "lock".
	slowReaderEliminationStrategy SlowReaderEliminationStrategy
	broadcastOnConnectionLoss bool
	broadcastChannels *list.List

	listenRequestch chan listenRequest

	lock sync.Mutex
	channels map[string] *listenSet
	closed bool
	// provide an escape hatch for goroutines sending on listenRequestch
	closeChannel chan bool
}

// NewNotifyDispatcher creates a new NotifyDispatcher, using the supplied
// pq.Listener underneath.  The ownership of the Listener is transferred to
// NotifyDispatcher.  You should not use it after calling NewNotifyDispatcher.
func NewNotifyDispatcher(l *pq.Listener) *NotifyDispatcher {
	d := &NotifyDispatcher{
		listener: l,
		slowReaderEliminationStrategy: CloseSlowReaders,
		broadcastOnConnectionLoss: true,
		broadcastChannels: list.New(),
		listenRequestch: make(chan listenRequest, 64),
		channels: make(map[string] *listenSet),
		closeChannel: make(chan bool),
	}
	go d.dispatcherLoop()
	go d.listenRequestHandlerLoop()
	return d
}

// Sets the strategy for mitigating the adverse effects slow readers might have
// on the dispatcher.  See SlowReaderEliminationStrategy.
func (d *NotifyDispatcher) SetSlowReaderEliminationStrategy(strategy SlowReaderEliminationStrategy) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.slowReaderEliminationStrategy = strategy
}

// Controls whether a nil notification from the underlying Listener is
// broadcast to all channels in the set.
func (d *NotifyDispatcher) SetBroadcastOnConnectionLoss(value bool) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.broadcastOnConnectionLoss = value
}

// Opens a new "broadcast channel".  A broadcast channel is sent to by the
// NotifyDispatcher every time the underlying Listener has re-established its
// database connection.
func (d *NotifyDispatcher) OpenBroadcastChannel() BroadcastChannel {
	d.lock.Lock()
	defer d.lock.Unlock()

	ch := BroadcastChannel{
		Channel: make(chan struct{}, 1),
	}
	ch.elem = d.broadcastChannels.PushFront(ch)
	return ch
}

// Closes the broadcast channel ch.
func (d *NotifyDispatcher) CloseBroadcastChannel(ch BroadcastChannel) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.broadcastChannels.Remove(ch.elem) != ch {
		panic("oops")
	}
	close(ch.Channel)
}

// Broadcast on all broadcastChannels and on all channels unless
// broadcastOnConnectionLoss is disabled.
func (d *NotifyDispatcher) broadcast() {
	d.lock.Lock()
	for e := d.broadcastChannels.Front(); e != nil; e = e.Next() {
		select {
			case e.Value.(BroadcastChannel).Channel <- struct{}{}:
			default:
				// nothing to do
		}
	}

	if !d.broadcastOnConnectionLoss {
		d.lock.Unlock()
		return
	}

	var reapchans []string
	for channel, set := range d.channels {
		if !set.broadcast(d.slowReaderEliminationStrategy, nil) {
			reapchans = append(reapchans, channel)
		}
	}
	d.lock.Unlock()

	for _, ch := range reapchans {
		d.listenRequestch <- listenRequest{ch, true}
	}
}

func (d *NotifyDispatcher) dispatch(n *pq.Notification) {
	reap := false

	d.lock.Lock()
	set, ok := d.channels[n.Channel]
	if ok {
		reap = !set.broadcast(d.slowReaderEliminationStrategy, n)
	}
	d.lock.Unlock()

	if reap {
		d.listenRequestch <- listenRequest{n.Channel, true}
	}
}

func (d *NotifyDispatcher) dispatcherLoop() {
	for {
		n := <-d.listener.Notify
		if n == nil {
			d.broadcast()
		} else {
			d.dispatch(n)
		}
	}
}

// Attempt to start listening on channel.  The caller should not be holding
// lock.
func (d *NotifyDispatcher) execListen(channel string) {
	for {
		err := d.listener.Listen(channel)
		// ErrChannelAlreadyOpen is a valid return value here; we could have
		// abandoned a channel in Unlisten() if the server returned an error
		// for no apparent reason.
		if err == nil ||
		   err == pq.ErrChannelAlreadyOpen {
			break
		}
	}

	d.lock.Lock()
	defer d.lock.Unlock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	set.setState(listenSetStateActive)
}

func (d *NotifyDispatcher) execUnlisten(channel string) {
	// we don't really care about the error
	_ = d.listener.Unlisten(channel)

	d.lock.Lock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	if set.state != listenSetStateZombie {
		panic("oops")
	}
	if set.reap() {
		delete(d.channels, channel)
		d.lock.Unlock()
	} else {
		// Couldn't reap the set because it got new listeners while we were
		// waiting for the UNLISTEN to go through.  Re-LISTEN it, but remember
		// to release the lock first.
		d.lock.Unlock()

		d.execListen(channel)
	}
}

func (d *NotifyDispatcher) listenRequestHandlerLoop() {
	for {
		// check closeChannel, just in case we've been closed and there's a
		// backlog of requests
		select {
			case <-d.closeChannel:
				return
			default:
		}

		select {
			case <-d.closeChannel:
				return

			case req := <-d.listenRequestch:
				if req.unlisten {
					d.execUnlisten(req.channel)
				} else {
					d.execListen(req.channel)
				}
		}
	}
}

func (d *NotifyDispatcher) requestListen(channel string, unlisten bool) error {
	select {
		// make sure we don't get stuck here if someone Close()s us
		case <-d.closeChannel:
			return errClosed
		case d.listenRequestch <- listenRequest{channel, unlisten}:
			return nil
	}

	panic("not reached")
}

// Listen adds ch to the set of listeners for notification channel channel.  ch
// should be a buffered channel.  If ch is already in the set of listeners for
// channel, ErrChannelAlreadyActive is returned.  After Listen has returned,
// the notification channel is open and the dispatcher will attempt to deliver
// all notifications received for that channel to ch.
//
// If SlowReaderEliminationStrategy is CloseSlowReaders, reusing the same ch
// for multiple notification channels is not allowed.
func (d *NotifyDispatcher) Listen(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	set, ok := d.channels[channel]
	if ok {
		err := set.add(ch)
		if err != nil {
			return err
		}
	} else {
		set = d.newListenSet(ch)
		d.channels[channel] = set

		// must not be holding the lock while requesting a listen
		d.lock.Unlock()
		err := d.requestListen(channel, false)
		if err != nil {
			return err
		}
		d.lock.Lock()
	}

	return set.waitForActive(&d.closed)
}

// Removes ch from the set of listeners for notification channel channel.  If
// ch is not in the set of listeners for channel, ErrChannelNotActive is
// returned.
func (d *NotifyDispatcher) Unlisten(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()

	if d.closed {
		d.lock.Unlock()
		return errClosed
	}

	set, ok := d.channels[channel]
	if !ok {
		d.lock.Unlock()
		return ErrChannelNotActive
	}
	last, err := set.remove(ch)
	d.lock.Unlock()
	if err != nil {
		return err
	}
	if !last {
		// the set isn't empty; nothing further for us to do
		return nil
	}

	// we were the last listener, cue the reaper
	return d.requestListen(channel, true)
}

func (d *NotifyDispatcher) Close() error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	d.closed = true
	close(d.closeChannel)

	for _, set := range d.channels {
		set.activeOrClosedCond.Broadcast()
	}

	// TODO: kill the dispatcher as well

	return nil
}

type listenSetState int
const (
	// The set was recently spawned or respawned, and it's waiting for a call
	// to Listen() to succeed.
	listenSetStateNewborn listenSetState = iota
	// The set is ready and any notifications from the database will be
	// dispatched to the set.
	listenSetStateActive
	// The set has recently been emptied, and it's waiting for a call to
	// Unlisten() to finish.
	listenSetStateZombie
)

type listenSet struct {
	channels map[chan<- *pq.Notification] struct{}
	state listenSetState
	activeOrClosedCond *sync.Cond
}

func (d *NotifyDispatcher) newListenSet(firstInhabitant chan<- *pq.Notification) *listenSet {
	s := &listenSet{
		channels: make(map[chan<- *pq.Notification] struct{}),
		state: listenSetStateNewborn,
	}
	s.activeOrClosedCond = sync.NewCond(&d.lock)
	s.channels[firstInhabitant] = struct{}{}
	return s
}

func (s *listenSet) setState(newState listenSetState) {
	var expectedState listenSetState
	switch newState {
		case listenSetStateNewborn:
			expectedState = listenSetStateZombie
		case listenSetStateActive:
			expectedState = listenSetStateNewborn
		case listenSetStateZombie:
			expectedState = listenSetStateActive
	}
	if s.state != expectedState {
		panic(fmt.Sprintf("illegal state transition from %v to %v", s.state, newState))
	}
	s.state = newState
	if s.state == listenSetStateActive {
		s.activeOrClosedCond.Broadcast()
	}
}

func (s *listenSet) add(ch chan<- *pq.Notification) error {
	_, ok := s.channels[ch]
	if ok {
		return ErrChannelAlreadyActive
	}
	s.channels[ch] = struct{}{}
	return nil
}

// Removes ch from the listen set.  last is false if ch was the last listener
// in the set and the caller should request an UNLISTEN.
func (s *listenSet) remove(ch chan<- *pq.Notification) (last bool, err error) {
	_, ok := s.channels[ch]
	if !ok {
		return false, ErrChannelNotActive
	}
	delete(s.channels, ch)

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return true, nil
	}
	return false, nil
}

// Sends n to all listeners in the set, using the supplied
// SlowReaderEliminationStrategy.  n may be nil.  Returns false if the set is
// empty (or was emptied as a result of eliminating slow readers) and the
// caller should request an UNLISTEN on it.  The caller should be holding
// d.lock.
func (s *listenSet) broadcast(strategy SlowReaderEliminationStrategy, n *pq.Notification) bool {
	// must be active
	if s.state != listenSetStateActive {
		return true
	}

	for ch := range s.channels {
		select {
			case ch <- n:

			default:
				if strategy == CloseSlowReaders {
					delete(s.channels, ch)
					close(ch)
				}
		}
	}

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return false
	}

	return true
}

// Marks the set active after a successful call to Listen().
func (s *listenSet) markActive() {
	s.setState(listenSetStateActive)
}

// Wait for the listen set to become "active".  Returns nil if successfull, or
// errClosed if the dispatcher was closed while waiting.  The caller should be
// holding d.lock.
func (s *listenSet) waitForActive(closed *bool) error {
	for {
		if *closed {
			return errClosed
		}
		if s.state == listenSetStateActive {
			return nil
		}
		s.activeOrClosedCond.Wait()
	}
}

// Try to reap a zombie set after Unlisten().  Returns true if the set should
// be removed, false otherwise.
func (s *listenSet) reap() bool {
	if s.state != listenSetStateZombie {
		panic("unexpected state in reap")
	}

	if len(s.channels) > 0 {
		// we need to be respawned
		s.setState(listenSetStateNewborn)
		return false
	}

	return true
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

// TODO: consider moving it to a more generic package.
package util

import (
	"container/heap"
	"time"
)

const (
	// TODO: Investigate what capacity is right.
	delayingDelivererUpdateChanCapacity = 1000
)

// DelayingDelivererItem is structure delivered by DelayingDeliverer to the
// target channel.
type DelayingDelivererItem struct {
	// Key under which the value was added to deliverer.
	Key string
	// Value of the item.
	Value interface{}
	// When the item should be delivered.
	DeliveryTime time.Time
}

type delivererHeap struct {
	keyPosition map[string]int
	data        []*DelayingDelivererItem
}

// Functions required by container.Heap.

func (dh *delivererHeap) Len() int { return len(dh.data) }
func (dh *delivererHeap) Less(i, j int) bool {
	return dh.data[i].DeliveryTime.Before(dh.data[j].DeliveryTime)
}
func (dh *delivererHeap) Swap(i, j int) {
	dh.keyPosition[dh.data[i].Key] = j
	dh.keyPosition[dh.data[j].Key] = i
	dh.data[i], dh.data[j] = dh.data[j], dh.data[i]
}

func (dh *delivererHeap) Push(x interface{}) {
	item := x.(*DelayingDelivererItem)
	dh.data = append(dh.data, item)
	dh.keyPosition[item.Key] = len(dh.data) - 1
}

func (dh *delivererHeap) Pop() interface{} {
	n := len(dh.data)
	item := dh.data[n-1]
	dh.data = dh.data[:n-1]
	delete(dh.keyPosition, item.Key)
	return item
}

// A structure that pushes the items to the target channel at a given time.
type DelayingDeliverer struct {
	// Channel to deliver the data when their time comes.
	targetChannel chan *DelayingDelivererItem
	// Store for data
	heap *delivererHeap
	// Channel to feed the main goroutine with updates.
	updateChannel chan *DelayingDelivererItem
	// To stop the main goroutine.
	stopChannel chan struct{}
}

func NewDelayingDeliverer(targetChannel chan *DelayingDelivererItem) *DelayingDeliverer {
	d := &DelayingDeliverer{
		targetChannel: targetChannel,
		heap: &delivererHeap{
			keyPosition: make(map[string]int),
			data:        make([]*DelayingDelivererItem, 0),
		},
		updateChannel: make(chan *DelayingDelivererItem, delayingDelivererUpdateChanCapacity),
		stopChannel:   make(chan struct{}),
	}
	go d.run()
	return d
}

// Deliver all items due before or equal to timestamp.
func (d *DelayingDeliverer) deliver(timestamp time.Time) {
	for d.heap.Len() > 0 {
		if timestamp.Before(d.heap.data[0].DeliveryTime) {
			return
		}
		item := heap.Pop(d.heap).(*DelayingDelivererItem)
		d.targetChannel <- item
	}
}

func (d *DelayingDeliverer) run() {
	for {
		now := time.Now()
		d.deliver(now)

		nextWakeUp := now.Add(time.Hour)
		if d.heap.Len() > 0 {
			nextWakeUp = d.heap.data[0].DeliveryTime
		}
		sleepTime := nextWakeUp.Sub(now)

		select {
		case <-time.After(sleepTime):
			break // just wake up and process the data
		case item := <-d.updateChannel:
			if position, found := d.heap.keyPosition[item.Key]; found {
				if item.DeliveryTime.Before(d.heap.data[position].DeliveryTime) {
					d.heap.data[position] = item
					heap.Fix(d.heap, position)
				}
				// Ignore if later.
			} else {
				heap.Push(d.heap, item)
			}
		case <-d.stopChannel:
			return
		}
	}
}

func (d *DelayingDeliverer) Stop() {
	close(d.stopChannel)
}

func (d *DelayingDeliverer) DeliverAt(key string, value interface{}, deliveryTime time.Time) {
	d.updateChannel <- &DelayingDelivererItem{
		Key:          key,
		Value:        value,
		DeliveryTime: deliveryTime,
	}
}
<commit_msg>Federation common libs - deliverAfer for delaying deliverer<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

// TODO: consider moving it to a more generic package.
package util

import (
	"container/heap"
	"time"
)

const (
	// TODO: Investigate what capacity is right.
	delayingDelivererUpdateChanCapacity = 1000
)

// DelayingDelivererItem is structure delivered by DelayingDeliverer to the
// target channel.
type DelayingDelivererItem struct {
	// Key under which the value was added to deliverer.
	Key string
	// Value of the item.
	Value interface{}
	// When the item should be delivered.
	DeliveryTime time.Time
}

type delivererHeap struct {
	keyPosition map[string]int
	data        []*DelayingDelivererItem
}

// Functions required by container.Heap.

func (dh *delivererHeap) Len() int { return len(dh.data) }
func (dh *delivererHeap) Less(i, j int) bool {
	return dh.data[i].DeliveryTime.Before(dh.data[j].DeliveryTime)
}
func (dh *delivererHeap) Swap(i, j int) {
	dh.keyPosition[dh.data[i].Key] = j
	dh.keyPosition[dh.data[j].Key] = i
	dh.data[i], dh.data[j] = dh.data[j], dh.data[i]
}

func (dh *delivererHeap) Push(x interface{}) {
	item := x.(*DelayingDelivererItem)
	dh.data = append(dh.data, item)
	dh.keyPosition[item.Key] = len(dh.data) - 1
}

func (dh *delivererHeap) Pop() interface{} {
	n := len(dh.data)
	item := dh.data[n-1]
	dh.data = dh.data[:n-1]
	delete(dh.keyPosition, item.Key)
	return item
}

// A structure that pushes the items to the target channel at a given time.
type DelayingDeliverer struct {
	// Channel to deliver the data when their time comes.
	targetChannel chan *DelayingDelivererItem
	// Store for data
	heap *delivererHeap
	// Channel to feed the main goroutine with updates.
	updateChannel chan *DelayingDelivererItem
	// To stop the main goroutine.
	stopChannel chan struct{}
}

func NewDelayingDeliverer(targetChannel chan *DelayingDelivererItem) *DelayingDeliverer {
	d := &DelayingDeliverer{
		targetChannel: targetChannel,
		heap: &delivererHeap{
			keyPosition: make(map[string]int),
			data:        make([]*DelayingDelivererItem, 0),
		},
		updateChannel: make(chan *DelayingDelivererItem, delayingDelivererUpdateChanCapacity),
		stopChannel:   make(chan struct{}),
	}
	go d.run()
	return d
}

// Deliver all items due before or equal to timestamp.
func (d *DelayingDeliverer) deliver(timestamp time.Time) {
	for d.heap.Len() > 0 {
		if timestamp.Before(d.heap.data[0].DeliveryTime) {
			return
		}
		item := heap.Pop(d.heap).(*DelayingDelivererItem)
		d.targetChannel <- item
	}
}

func (d *DelayingDeliverer) run() {
	for {
		now := time.Now()
		d.deliver(now)

		nextWakeUp := now.Add(time.Hour)
		if d.heap.Len() > 0 {
			nextWakeUp = d.heap.data[0].DeliveryTime
		}
		sleepTime := nextWakeUp.Sub(now)

		select {
		case <-time.After(sleepTime):
			break // just wake up and process the data
		case item := <-d.updateChannel:
			if position, found := d.heap.keyPosition[item.Key]; found {
				if item.DeliveryTime.Before(d.heap.data[position].DeliveryTime) {
					d.heap.data[position] = item
					heap.Fix(d.heap, position)
				}
				// Ignore if later.
			} else {
				heap.Push(d.heap, item)
			}
		case <-d.stopChannel:
			return
		}
	}
}

func (d *DelayingDeliverer) Stop() {
	close(d.stopChannel)
}

// Delivers value at the given time.
func (d *DelayingDeliverer) DeliverAt(key string, value interface{}, deliveryTime time.Time) {
	d.updateChannel <- &DelayingDelivererItem{
		Key:          key,
		Value:        value,
		DeliveryTime: deliveryTime,
	}
}

// Delivers value after the given delay.
func (d *DelayingDeliverer) DeliverAfter(key string, value interface{}, delay time.Duration) {
	d.DeliverAt(key, value, time.Now().Add(delay))
}
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 obcca

import (
	"testing"
	//"encoding/json"
	//"errors"
	"fmt"
	//"io/ioutil"
	"net"
	"os"
	//"runtime"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"github.com/spf13/viper"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"

	//"github.com/spf13/cobra"
	//"github.com/openblockchain/obc-peer/events/producer"
	"github.com/openblockchain/obc-peer/openchain"
	"github.com/openblockchain/obc-peer/openchain/chaincode"
	"github.com/openblockchain/obc-peer/openchain/consensus/helper"
	"github.com/openblockchain/obc-peer/openchain/ledger/genesis"
	"github.com/openblockchain/obc-peer/openchain/peer"
	"github.com/openblockchain/obc-peer/openchain/rest"
	pb "github.com/openblockchain/obc-peer/protos"
	"github.com/openblockchain/obc-peer/openchain/util"

	"sync"
	"io/ioutil"
)

func TestMain(m *testing.M) {
	setupTestConfig()
	os.Exit(m.Run())
}

func setupTestConfig() {
	viper.AutomaticEnv()
	viper.SetConfigName("obcca_test") // name of config file (without extension)
	viper.AddConfigPath("./")        // path to look for the config file in
	viper.AddConfigPath("./..")      // path to look for the config file in
	viper.AddConfigPath("./../..")      // path to look for the config file in
	err := viper.ReadInConfig()      // Find and read the config file
	if err != nil {                  // Handle errors reading the config file
		panic(fmt.Errorf("Fatal error config file: %s \n", err))
	}

	chaincodeDevMode := false
	if chaincodeDevMode {
		logger.Info("Running in chaincode development mode. Set consensus to NOOPS and user starts chaincode")
		viper.Set("peer.validator.enabled", true)
		viper.Set("peer.validator.consensus", "noops")
		viper.Set("chaincode.mode", chaincode.DevModeUserRunsChaincode)
	}

	/*viper.Set("peer.fileSystemPath", "/var/openchain/test")
	viper.Set("security.enabled", true)
	viper.Set("ports.ecaP", ":50051")
	viper.Set("ports.ecaA", ":50052")
	viper.Set("ports.tcaP", ":50551")
	viper.Set("ports.tcaA", ":50552")
	viper.Set("ports.tlscaP", ":50951")
	viper.Set("ports.tlscaA", ":50952")
	viper.Set("hosts.eca", "localhost")
	viper.Set("hosts.tca", "localhost")
	viper.Set("hosts.tlsca", "localhost")
	viper.Set("eca.users.nepumuk", "9gvZQRwhUq9q")
	viper.Set("eca.users.jim", "AwbeJH2kw9qK")
	viper.Set("eca.users.lukas", "NPKYL39uKbkj")
	viper.Set("eca.users.system_chaincode_invoker", "DRJ20pEql15a")
	viper.Set("pki.validity-period.update", true)
	viper.Set("pki.validity-period.tls.enabled", false)
	viper.Set("pki.validity-period.tls.cert.file", "testdata/server1.pem")
	viper.Set("pki.validity-period.tls.key.file", "testdata/server1.key")
	viper.Set("pki.validity-period.devops-address", "0.0.0.0:30303")*/
}

func TestValidityPeriod(t *testing.T) {

	startTCA()

	err := startOpenchain()
	if(err != nil){
		t.Logf("Error starting Openchain: %s", err)
		t.Fail()
	}
}

func startTCA() {
	LogInit(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr, os.Stdout)

	eca := NewECA()
	defer eca.Close()

	tca := NewTCA(eca)
	defer tca.Close()

	var wg sync.WaitGroup
	eca.Start(&wg)
	tca.Start(&wg)

	//exampleQueryTransaction(context.Background(),"github.com/openblockchain/obc-peer/openchain/system_chaincode/validity_period_update", "0.0.1", []string{})

	//wg.Wait()
}

// getChaincodeID constructs the ID from pb.ChaincodeID; used by handlerMap
func getChaincodeID(cID *pb.ChaincodeID) (string, error) {
	if cID == nil {
		return "", fmt.Errorf("Cannot construct chaincodeID, got nil object")
	}
	var urlLocation string
	if strings.HasPrefix(cID.Url, "http://") {
		urlLocation = cID.Url[7:]
	} else if strings.HasPrefix(cID.Url, "https://") {
		urlLocation = cID.Url[8:]
	} else {
		urlLocation = cID.Url
	}
	return urlLocation + ":" + cID.Version, nil
}

func exampleQueryTransaction(ctxt context.Context, url string, version string, args []string) error {

	chaincodeID, _ := getChaincodeID(&pb.ChaincodeID{Url: url, Version: version})

	fmt.Printf("Going to query\n")
	f := "query"
	spec := &pb.ChaincodeSpec{Type: pb.ChaincodeSpec_GOLANG,
		ChaincodeID: &pb.ChaincodeID{Url: chaincodePath,
			Version: chaincodeVersion,
		},
		CtorMsg: &pb.ChaincodeInput{Function: f, Args: args},
		}

	uuid, _, err := invoke(ctxt, spec, pb.Transaction_CHAINCODE_QUERY)

	fmt.Println(uuid)
	if err != nil {
		return fmt.Errorf("Error querying <%s>: %s", chaincodeID, err)
	}

	return nil
}

// Invoke or query a chaincode.
func invoke(ctx context.Context, spec *pb.ChaincodeSpec, typ pb.Transaction_Type) (string, []byte, error) {
	chaincodeInvocationSpec := &pb.ChaincodeInvocationSpec{ChaincodeSpec: spec}

	// Now create the Transactions message and send to Peer.
	uuid, uuidErr := util.GenerateUUID()
	if uuidErr != nil {
		return "", nil, uuidErr
	}
	transaction, err := pb.NewChaincodeExecute(chaincodeInvocationSpec, uuid, typ)
	if err != nil {
		return uuid, nil, fmt.Errorf("Error invoking chaincode: %s ", err)
	}
	retval, err := chaincode.Execute(ctx, chaincode.GetChain("default"), transaction, nil)
	return uuid, retval, err
}

func startOpenchain() error {

	peerEndpoint, err := peer.GetPeerEndpoint()
	if err != nil {
		logger.Error(fmt.Sprintf("Failed to get Peer Endpoint: %s", err))
		return err
	}

	listenAddr := viper.GetString("peer.listenaddress")

	if "" == listenAddr {
		logger.Debug("Listen address not specified, using peer endpoint address")
		listenAddr = peerEndpoint.Address
	}

	lis, err := net.Listen("tcp", listenAddr)
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}

//	ehubLis, ehubGrpcServer, err := createEventHubServer()
//	if err != nil {
//		grpclog.Fatalf("failed to create ehub server: %v", err)
//	}

	logger.Info("Security enabled status: %t", viper.GetBool("security.enabled"))

	var opts []grpc.ServerOption
	if viper.GetBool("peer.tls.enabled") {
		creds, err := credentials.NewServerTLSFromFile(viper.GetString("peer.tls.cert.file"), viper.GetString("peer.tls.key.file"))
		if err != nil {
			grpclog.Fatalf("Failed to generate credentials %v", err)
		}
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}

	grpcServer := grpc.NewServer(opts...)

	// Register the Peer server
	//pb.RegisterPeerServer(grpcServer, openchain.NewPeer())
	var peerServer *peer.PeerImpl

	if viper.GetBool("peer.validator.enabled") {
		logger.Debug("Running as validating peer - installing consensus %s", viper.GetString("peer.validator.consensus"))
		peerServer, _ = peer.NewPeerWithHandler(helper.NewConsensusHandler)
	} else {
		logger.Debug("Running as non-validating peer")
		peerServer, _ = peer.NewPeerWithHandler(peer.NewPeerHandler)
	}
	pb.RegisterPeerServer(grpcServer, peerServer)

	// Register the Admin server
	pb.RegisterAdminServer(grpcServer, openchain.NewAdminServer())

	// Register ChaincodeSupport server...
	// TODO : not the "DefaultChain" ... we have to revisit when we do multichain
	registerChaincodeSupport(chaincode.DefaultChain, grpcServer)

	// Register Devops server
	serverDevops := openchain.NewDevopsServer(peerServer)
	pb.RegisterDevopsServer(grpcServer, serverDevops)

	// Register the ServerOpenchain server
	serverOpenchain, err := openchain.NewOpenchainServer()
	if err != nil {
		logger.Error(fmt.Sprintf("Error creating OpenchainServer: %s", err))
		return err
	}

	pb.RegisterOpenchainServer(grpcServer, serverOpenchain)

	// Create and register the REST service
	go rest.StartOpenchainRESTServer(serverOpenchain, serverDevops)

	rootNode, err := openchain.GetRootNode()
	if err != nil {
		grpclog.Fatalf("Failed to get peer.discovery.rootnode valey: %s", err)
	}

	logger.Info("Starting peer with id=%s, network id=%s, address=%s, discovery.rootnode=%s, validator=%v",
		peerEndpoint.ID, viper.GetString("peer.networkId"),
		peerEndpoint.Address, rootNode, viper.GetBool("peer.validator.enabled"))

	// Start the grpc server. Done in a goroutine so we can deploy the
	// genesis block if needed.
	serve := make(chan bool)
	go func() {
		grpcServer.Serve(lis)
		serve <- true
	}()

	// Deploy the geneis block if needed.
	if viper.GetBool("peer.validator.enabled") {
		makeGeneisError := genesis.MakeGenesis(peerServer.GetSecHelper())
		if makeGeneisError != nil {
			return makeGeneisError
		}
	}

	//start the event hub server
//	if ehubGrpcServer != nil && ehubLis != nil {
//		go ehubGrpcServer.Serve(ehubLis)
//	}

	// Block until grpc server exits
	<-serve

	return nil
}

func registerChaincodeSupport(chainname chaincode.ChainName, grpcServer *grpc.Server) {
	//get user mode
	userRunsCC := false
	if viper.GetString("chaincode.mode") == chaincode.DevModeUserRunsChaincode {
		userRunsCC = true
	}

	//get chaincode startup timeout
	tOut, err := strconv.Atoi(viper.GetString("chaincode.startuptimeout"))
	if err != nil { //what went wrong ?
		fmt.Printf("could not retrive timeout var...setting to 5secs\n")
		tOut = 5000
	}
	ccStartupTimeout := time.Duration(tOut) * time.Millisecond

	pb.RegisterChaincodeSupportServer(grpcServer, chaincode.NewChaincodeSupport(chainname, peer.GetPeerEndpoint, userRunsCC, ccStartupTimeout))
}<commit_msg>Query function<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 obcca

import (
	"testing"
	//"encoding/json"
	//"errors"
	"fmt"
	//"io/ioutil"
	"net"
	"os"
	//"runtime"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"github.com/spf13/viper"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"

	//"github.com/spf13/cobra"
	//"github.com/openblockchain/obc-peer/events/producer"
	"github.com/openblockchain/obc-peer/openchain"
	"github.com/openblockchain/obc-peer/openchain/chaincode"
	"github.com/openblockchain/obc-peer/openchain/consensus/helper"
	"github.com/openblockchain/obc-peer/openchain/ledger/genesis"
	"github.com/openblockchain/obc-peer/openchain/peer"
	"github.com/openblockchain/obc-peer/openchain/rest"
	pb "github.com/openblockchain/obc-peer/protos"
	"github.com/openblockchain/obc-peer/openchain/util"

	"sync"
	"io/ioutil"
)

func TestMain(m *testing.M) {
	setupTestConfig()
	os.Exit(m.Run())
}

func setupTestConfig() {
	viper.AutomaticEnv()
	viper.SetConfigName("obcca_test") // name of config file (without extension)
	viper.AddConfigPath("./")        // path to look for the config file in
	viper.AddConfigPath("./..")      // path to look for the config file in
	viper.AddConfigPath("./../..")      // path to look for the config file in
	err := viper.ReadInConfig()      // Find and read the config file
	if err != nil {                  // Handle errors reading the config file
		panic(fmt.Errorf("Fatal error config file: %s \n", err))
	}

	chaincodeDevMode := false
	if chaincodeDevMode {
		logger.Info("Running in chaincode development mode. Set consensus to NOOPS and user starts chaincode")
		viper.Set("peer.validator.enabled", true)
		viper.Set("peer.validator.consensus", "noops")
		viper.Set("chaincode.mode", chaincode.DevModeUserRunsChaincode)
	}

	/*viper.Set("peer.fileSystemPath", "/var/openchain/test")
	viper.Set("security.enabled", true)
	viper.Set("ports.ecaP", ":50051")
	viper.Set("ports.ecaA", ":50052")
	viper.Set("ports.tcaP", ":50551")
	viper.Set("ports.tcaA", ":50552")
	viper.Set("ports.tlscaP", ":50951")
	viper.Set("ports.tlscaA", ":50952")
	viper.Set("hosts.eca", "localhost")
	viper.Set("hosts.tca", "localhost")
	viper.Set("hosts.tlsca", "localhost")
	viper.Set("eca.users.nepumuk", "9gvZQRwhUq9q")
	viper.Set("eca.users.jim", "AwbeJH2kw9qK")
	viper.Set("eca.users.lukas", "NPKYL39uKbkj")
	viper.Set("eca.users.system_chaincode_invoker", "DRJ20pEql15a")
	viper.Set("pki.validity-period.update", true)
	viper.Set("pki.validity-period.tls.enabled", false)
	viper.Set("pki.validity-period.tls.cert.file", "testdata/server1.pem")
	viper.Set("pki.validity-period.tls.key.file", "testdata/server1.key")
	viper.Set("pki.validity-period.devops-address", "0.0.0.0:30303")*/
}

func TestValidityPeriod(t *testing.T) {

	startTCA()

	err := startOpenchain()
	if(err != nil){
		t.Logf("Error starting Openchain: %s", err)
		t.Fail()
	}
}

func startTCA() {
	LogInit(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr, os.Stdout)

	eca := NewECA()
	defer eca.Close()

	tca := NewTCA(eca)
	defer tca.Close()

	var wg sync.WaitGroup
	eca.Start(&wg)
	tca.Start(&wg)

	exampleQueryTransaction(context.Background(),"github.com/openblockchain/obc-peer/openchain/system_chaincode/validity_period_update", "0.0.1", []string{"system.validity.period"})

	//wg.Wait()
}

// getChaincodeID constructs the ID from pb.ChaincodeID; used by handlerMap
func getChaincodeID(cID *pb.ChaincodeID) (string, error) {
	if cID == nil {
		return "", fmt.Errorf("Cannot construct chaincodeID, got nil object")
	}
	var urlLocation string
	if strings.HasPrefix(cID.Url, "http://") {
		urlLocation = cID.Url[7:]
	} else if strings.HasPrefix(cID.Url, "https://") {
		urlLocation = cID.Url[8:]
	} else {
		urlLocation = cID.Url
	}
	return urlLocation + ":" + cID.Version, nil
}

func exampleQueryTransaction(ctxt context.Context, url string, version string, args []string) error {

	chaincodeID, _ := getChaincodeID(&pb.ChaincodeID{Url: url, Version: version})

	fmt.Printf("Going to query\n")
	f := "query"
	spec := &pb.ChaincodeSpec{Type: pb.ChaincodeSpec_GOLANG,
		ChaincodeID: &pb.ChaincodeID{Url: chaincodePath,
			Version: chaincodeVersion,
		},
		CtorMsg: &pb.ChaincodeInput{Function: f, Args: args},
		}

	spec.SecureContext = "system_chaincode_invoker"
	uuid, _, err := invoke(ctxt, spec, pb.Transaction_CHAINCODE_QUERY)

	fmt.Println(uuid)
	if err != nil {
		return fmt.Errorf("Error querying <%s>: %s", chaincodeID, err)
	}

	return nil
}

func getDevopsClient(peerAddress string) (pb.DevopsClient, error) {
	var opts []grpc.DialOption
	if viper.GetBool("pki.validity-period.tls.enabled") {
		var sn string
		if viper.GetString("pki.validity-period.tls.server-host-override") != "" {
			sn = viper.GetString("pki.validity-period.tls.server-host-override")
		}
		var creds credentials.TransportAuthenticator
		if viper.GetString("pki.validity-period.tls.cert.file") != "" {
			var err error
			creds, err = credentials.NewClientTLSFromFile(viper.GetString("pki.validity-period.tls.cert.file"), sn)
			if err != nil {
				grpclog.Fatalf("Failed to create TLS credentials %v", err)
			}
		} else {
			creds = credentials.NewClientTLSFromCert(nil, sn)
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	}
	opts = append(opts, grpc.WithTimeout(systemChaincodeTimeout))
	opts = append(opts, grpc.WithBlock())
	opts = append(opts, grpc.WithInsecure())
	conn, err := grpc.Dial(peerAddress, opts...)

	if err != nil {
		return nil, fmt.Errorf("Error trying to connect to local peer: %s", err)
	}

	devopsClient := obc.NewDevopsClient(conn)
	return devopsClient, nil
}

// Invoke or query a chaincode.
func invoke(ctx context.Context, spec *pb.ChaincodeSpec, typ pb.Transaction_Type) (string, []byte, error) {
	chaincodeInvocationSpec := &pb.ChaincodeInvocationSpec{ChaincodeSpec: spec}

	// Now create the Transactions message and send to Peer.
	uuid, uuidErr := util.GenerateUUID()
	if uuidErr != nil {
		return "", nil, uuidErr
	}
	transaction, err := pb.NewChaincodeExecute(chaincodeInvocationSpec, uuid, typ)
	if err != nil {
		return uuid, nil, fmt.Errorf("Error invoking chaincode: %s ", err)
	}
	retval, err := chaincode.Execute(ctx, chaincode.GetChain("default"), transaction, nil)
	return uuid, retval, err
}

func startOpenchain() error {

	peerEndpoint, err := peer.GetPeerEndpoint()
	if err != nil {
		logger.Error(fmt.Sprintf("Failed to get Peer Endpoint: %s", err))
		return err
	}

	listenAddr := viper.GetString("peer.listenaddress")

	if "" == listenAddr {
		logger.Debug("Listen address not specified, using peer endpoint address")
		listenAddr = peerEndpoint.Address
	}

	lis, err := net.Listen("tcp", listenAddr)
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}

//	ehubLis, ehubGrpcServer, err := createEventHubServer()
//	if err != nil {
//		grpclog.Fatalf("failed to create ehub server: %v", err)
//	}

	logger.Info("Security enabled status: %t", viper.GetBool("security.enabled"))

	var opts []grpc.ServerOption
	if viper.GetBool("peer.tls.enabled") {
		creds, err := credentials.NewServerTLSFromFile(viper.GetString("peer.tls.cert.file"), viper.GetString("peer.tls.key.file"))
		if err != nil {
			grpclog.Fatalf("Failed to generate credentials %v", err)
		}
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}

	grpcServer := grpc.NewServer(opts...)

	// Register the Peer server
	//pb.RegisterPeerServer(grpcServer, openchain.NewPeer())
	var peerServer *peer.PeerImpl

	if viper.GetBool("peer.validator.enabled") {
		logger.Debug("Running as validating peer - installing consensus %s", viper.GetString("peer.validator.consensus"))
		peerServer, _ = peer.NewPeerWithHandler(helper.NewConsensusHandler)
	} else {
		logger.Debug("Running as non-validating peer")
		peerServer, _ = peer.NewPeerWithHandler(peer.NewPeerHandler)
	}
	pb.RegisterPeerServer(grpcServer, peerServer)

	// Register the Admin server
	pb.RegisterAdminServer(grpcServer, openchain.NewAdminServer())

	// Register ChaincodeSupport server...
	// TODO : not the "DefaultChain" ... we have to revisit when we do multichain
	registerChaincodeSupport(chaincode.DefaultChain, grpcServer)

	// Register Devops server
	serverDevops := openchain.NewDevopsServer(peerServer)
	pb.RegisterDevopsServer(grpcServer, serverDevops)

	// Register the ServerOpenchain server
	serverOpenchain, err := openchain.NewOpenchainServer()
	if err != nil {
		logger.Error(fmt.Sprintf("Error creating OpenchainServer: %s", err))
		return err
	}

	pb.RegisterOpenchainServer(grpcServer, serverOpenchain)

	// Create and register the REST service
	go rest.StartOpenchainRESTServer(serverOpenchain, serverDevops)

	rootNode, err := openchain.GetRootNode()
	if err != nil {
		grpclog.Fatalf("Failed to get peer.discovery.rootnode valey: %s", err)
	}

	logger.Info("Starting peer with id=%s, network id=%s, address=%s, discovery.rootnode=%s, validator=%v",
		peerEndpoint.ID, viper.GetString("peer.networkId"),
		peerEndpoint.Address, rootNode, viper.GetBool("peer.validator.enabled"))

	// Start the grpc server. Done in a goroutine so we can deploy the
	// genesis block if needed.
	serve := make(chan bool)
	go func() {
		grpcServer.Serve(lis)
		serve <- true
	}()

	// Deploy the geneis block if needed.
	if viper.GetBool("peer.validator.enabled") {
		makeGeneisError := genesis.MakeGenesis(peerServer.GetSecHelper())
		if makeGeneisError != nil {
			return makeGeneisError
		}
	}

	//start the event hub server
//	if ehubGrpcServer != nil && ehubLis != nil {
//		go ehubGrpcServer.Serve(ehubLis)
//	}

	// Block until grpc server exits
	<-serve

	return nil
}

func registerChaincodeSupport(chainname chaincode.ChainName, grpcServer *grpc.Server) {
	//get user mode
	userRunsCC := false
	if viper.GetString("chaincode.mode") == chaincode.DevModeUserRunsChaincode {
		userRunsCC = true
	}

	//get chaincode startup timeout
	tOut, err := strconv.Atoi(viper.GetString("chaincode.startuptimeout"))
	if err != nil { //what went wrong ?
		fmt.Printf("could not retrive timeout var...setting to 5secs\n")
		tOut = 5000
	}
	ccStartupTimeout := time.Duration(tOut) * time.Millisecond

	pb.RegisterChaincodeSupportServer(grpcServer, chaincode.NewChaincodeSupport(chainname, peer.GetPeerEndpoint, userRunsCC, ccStartupTimeout))
}<|endoftext|>


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

package cmd

import (
	"bytes"
	"context"
	"encoding/hex"
	"hash"
	"io"

	"github.com/minio/minio/cmd/logger"
)

// Calculates bitrot in chunks and writes the hash into the stream.
type streamingBitrotWriter struct {
	iow       *io.PipeWriter
	h         hash.Hash
	shardSize int64
	canClose  chan struct{} // Needed to avoid race explained in Close() call.
}

func (b *streamingBitrotWriter) Write(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	b.h.Reset()
	b.h.Write(p)
	hashBytes := b.h.Sum(nil)
	n, err := b.iow.Write(hashBytes)
	if n != len(hashBytes) {
		logger.LogIf(context.Background(), err)
		return 0, err
	}
	n, err = b.iow.Write(p)
	return n, err
}

func (b *streamingBitrotWriter) Close() error {
	err := b.iow.Close()
	// Wait for all data to be written before returning else it causes race conditions.
	// Race condition is because of io.PipeWriter implementation. i.e consider the following
	// sequent of operations:
	// 1) pipe.Write()
	// 2) pipe.Close()
	// Now pipe.Close() can return before the data is read on the other end of the pipe and written to the disk
	// Hence an immediate Read() on the file can return incorrect data.
	<-b.canClose
	return err
}

// Returns streaming bitrot writer implementation.
func newStreamingBitrotWriter(disk StorageAPI, volume, filePath string, length int64, algo BitrotAlgorithm, shardSize int64) io.WriteCloser {
	r, w := io.Pipe()
	h := algo.New()
	bw := &streamingBitrotWriter{w, h, shardSize, make(chan struct{})}
	go func() {
		totalFileSize := int64(-1) // For compressed objects length will be unknown (represented by length=-1)
		if length != -1 {
			bitrotSumsTotalSize := ceilFrac(length, shardSize) * int64(h.Size()) // Size used for storing bitrot checksums.
			totalFileSize = bitrotSumsTotalSize + length
		}
		err := disk.CreateFile(volume, filePath, totalFileSize, r)
		if err != nil {
			reqInfo := (&logger.ReqInfo{}).AppendTags("storageDisk", disk.String())
			ctx := logger.SetReqInfo(context.Background(), reqInfo)
			logger.LogIf(ctx, err)
		}
		r.CloseWithError(err)
		close(bw.canClose)
	}()
	return bw
}

// ReadAt() implementation which verifies the bitrot hash available as part of the stream.
type streamingBitrotReader struct {
	disk       StorageAPI
	rc         io.ReadCloser
	volume     string
	filePath   string
	tillOffset int64
	currOffset int64
	h          hash.Hash
	shardSize  int64
	hashBytes  []byte
}

func (b *streamingBitrotReader) Close() error {
	if b.rc == nil {
		return nil
	}
	return b.rc.Close()
}

func (b *streamingBitrotReader) ReadAt(buf []byte, offset int64) (int, error) {
	var err error
	if offset%b.shardSize != 0 {
		// Offset should always be aligned to b.shardSize
		logger.LogIf(context.Background(), errUnexpected)
		return 0, errUnexpected
	}
	if b.rc == nil {
		// For the first ReadAt() call we need to open the stream for reading.
		b.currOffset = offset
		streamOffset := (offset/b.shardSize)*int64(b.h.Size()) + offset
		b.rc, err = b.disk.ReadFileStream(b.volume, b.filePath, streamOffset, b.tillOffset-streamOffset)
		if err != nil {
			reqInfo := (&logger.ReqInfo{}).AppendTags("storageDisk", b.disk.String())
			ctx := logger.SetReqInfo(context.Background(), reqInfo)
			logger.LogIf(ctx, err)
			return 0, err
		}
	}
	if offset != b.currOffset {
		logger.LogIf(context.Background(), errUnexpected)
		return 0, errUnexpected
	}
	b.h.Reset()
	_, err = io.ReadFull(b.rc, b.hashBytes)
	if err != nil {
		logger.LogIf(context.Background(), err)
		return 0, err
	}
	_, err = io.ReadFull(b.rc, buf)
	if err != nil {
		logger.LogIf(context.Background(), err)
		return 0, err
	}
	b.h.Write(buf)

	if !bytes.Equal(b.h.Sum(nil), b.hashBytes) {
		err = hashMismatchError{hex.EncodeToString(b.hashBytes), hex.EncodeToString(b.h.Sum(nil))}
		logger.LogIf(context.Background(), err)
		return 0, err
	}
	b.currOffset += int64(len(buf))
	return len(buf), nil
}

// Returns streaming bitrot reader implementation.
func newStreamingBitrotReader(disk StorageAPI, volume, filePath string, tillOffset int64, algo BitrotAlgorithm, shardSize int64) *streamingBitrotReader {
	h := algo.New()
	return &streamingBitrotReader{
		disk,
		nil,
		volume,
		filePath,
		ceilFrac(tillOffset, shardSize)*int64(h.Size()) + tillOffset,
		0,
		h,
		shardSize,
		make([]byte, h.Size()),
	}
}
<commit_msg>Remove logs from bitrot-streaming.go as erasure layer is already logging (#7603)<commit_after>/*
 * MinIO Cloud Storage, (C) 2019 MinIO, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cmd

import (
	"bytes"
	"context"
	"encoding/hex"
	"hash"
	"io"

	"github.com/minio/minio/cmd/logger"
)

// Calculates bitrot in chunks and writes the hash into the stream.
type streamingBitrotWriter struct {
	iow       *io.PipeWriter
	h         hash.Hash
	shardSize int64
	canClose  chan struct{} // Needed to avoid race explained in Close() call.
}

func (b *streamingBitrotWriter) Write(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	b.h.Reset()
	b.h.Write(p)
	hashBytes := b.h.Sum(nil)
	_, err := b.iow.Write(hashBytes)
	if err != nil {
		return 0, err
	}
	return b.iow.Write(p)
}

func (b *streamingBitrotWriter) Close() error {
	err := b.iow.Close()
	// Wait for all data to be written before returning else it causes race conditions.
	// Race condition is because of io.PipeWriter implementation. i.e consider the following
	// sequent of operations:
	// 1) pipe.Write()
	// 2) pipe.Close()
	// Now pipe.Close() can return before the data is read on the other end of the pipe and written to the disk
	// Hence an immediate Read() on the file can return incorrect data.
	<-b.canClose
	return err
}

// Returns streaming bitrot writer implementation.
func newStreamingBitrotWriter(disk StorageAPI, volume, filePath string, length int64, algo BitrotAlgorithm, shardSize int64) io.WriteCloser {
	r, w := io.Pipe()
	h := algo.New()
	bw := &streamingBitrotWriter{w, h, shardSize, make(chan struct{})}
	go func() {
		totalFileSize := int64(-1) // For compressed objects length will be unknown (represented by length=-1)
		if length != -1 {
			bitrotSumsTotalSize := ceilFrac(length, shardSize) * int64(h.Size()) // Size used for storing bitrot checksums.
			totalFileSize = bitrotSumsTotalSize + length
		}
		err := disk.CreateFile(volume, filePath, totalFileSize, r)
		r.CloseWithError(err)
		close(bw.canClose)
	}()
	return bw
}

// ReadAt() implementation which verifies the bitrot hash available as part of the stream.
type streamingBitrotReader struct {
	disk       StorageAPI
	rc         io.ReadCloser
	volume     string
	filePath   string
	tillOffset int64
	currOffset int64
	h          hash.Hash
	shardSize  int64
	hashBytes  []byte
}

func (b *streamingBitrotReader) Close() error {
	if b.rc == nil {
		return nil
	}
	return b.rc.Close()
}

func (b *streamingBitrotReader) ReadAt(buf []byte, offset int64) (int, error) {
	var err error
	if offset%b.shardSize != 0 {
		// Offset should always be aligned to b.shardSize
		// Can never happen unless there are programmer bugs
		return 0, errUnexpected
	}
	if b.rc == nil {
		// For the first ReadAt() call we need to open the stream for reading.
		b.currOffset = offset
		streamOffset := (offset/b.shardSize)*int64(b.h.Size()) + offset
		b.rc, err = b.disk.ReadFileStream(b.volume, b.filePath, streamOffset, b.tillOffset-streamOffset)
		if err != nil {
			return 0, err
		}
	}
	if offset != b.currOffset {
		// Can never happen unless there are programmer bugs
		return 0, errUnexpected
	}
	b.h.Reset()
	_, err = io.ReadFull(b.rc, b.hashBytes)
	if err != nil {
		return 0, err
	}
	_, err = io.ReadFull(b.rc, buf)
	if err != nil {
		return 0, err
	}
	b.h.Write(buf)

	if !bytes.Equal(b.h.Sum(nil), b.hashBytes) {
		err = hashMismatchError{hex.EncodeToString(b.hashBytes), hex.EncodeToString(b.h.Sum(nil))}
		logger.LogIf(context.Background(), err)
		return 0, err
	}
	b.currOffset += int64(len(buf))
	return len(buf), nil
}

// Returns streaming bitrot reader implementation.
func newStreamingBitrotReader(disk StorageAPI, volume, filePath string, tillOffset int64, algo BitrotAlgorithm, shardSize int64) *streamingBitrotReader {
	h := algo.New()
	return &streamingBitrotReader{
		disk,
		nil,
		volume,
		filePath,
		ceilFrac(tillOffset, shardSize)*int64(h.Size()) + tillOffset,
		0,
		h,
		shardSize,
		make([]byte, h.Size()),
	}
}
<|endoftext|>


<commit_before>package main

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

	"golang.org/x/crypto/ssh"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	"github.com/concourse/bin/bindata"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/restart"
	"github.com/tedsuo/ifrit/sigmon"
)

type WorkerCommand struct {
	Name string   `long:"name" description:"The name to set for the worker during registration. If not specified, the hostname will be used."`
	Tags []string `long:"tag" description:"A tag to set during registration. Can be specified multiple times."`

	TeamName string `long:"team" description:"The name of the team that this worker will be assigned to."`

	HTTPProxy  URLFlag  `long:"http-proxy"  env:"http_proxy"                  description:"HTTP proxy endpoint to use for containers."`
	HTTPSProxy URLFlag  `long:"https-proxy" env:"https_proxy"                 description:"HTTPS proxy endpoint to use for containers."`
	NoProxy    []string `long:"no-proxy"    env:"no_proxy"    env-delim:","   description:"Blacklist of addresses to skip the proxy when reaching."`

	WorkDir DirFlag `long:"work-dir" required:"true" description:"Directory in which to place container data."`

	BindIP   IPFlag `long:"bind-ip"   default:"127.0.0.1" description:"IP address on which to listen for the Garden server."`
	BindPort uint16 `long:"bind-port" default:"7777"      description:"Port on which to listen for the Garden server."`

	PeerIP IPFlag `long:"peer-ip" description:"IP used to reach this worker from the ATC nodes. If omitted, the worker will be forwarded through the SSH connection to the TSA."`

	Garden GardenBackend `group:"Garden Configuration" namespace:"garden"`

	Baggageclaim baggageclaimcmd.BaggageclaimCommand `group:"Baggageclaim Configuration" namespace:"baggageclaim"`

	TSA BeaconConfig `group:"TSA Configuration" namespace:"tsa"`

	Metrics struct {
		YellerAPIKey      string `long:"yeller-api-key"     description:"Yeller API key. If specified, all errors logged will be emitted."`
		YellerEnvironment string `long:"yeller-environment" description:"Environment to tag on all Yeller events emitted."`
	} `group:"Metrics & Diagnostics"`
}

func (cmd *WorkerCommand) Execute(args []string) error {
	logger := lager.NewLogger("worker")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.INFO))

	hasAssets, err := cmd.setup(logger.Session("setup"))
	if err != nil {
		return err
	}

	worker, gardenRunner, err := cmd.gardenRunner(logger.Session("garden"), hasAssets)
	if err != nil {
		return err
	}

	worker.Version = WorkerVersion

	baggageclaimRunner, err := cmd.baggageclaimRunner(logger.Session("baggageclaim"), hasAssets)
	if err != nil {
		return err
	}

	members := grouper.Members{
		{
			Name:   "garden",
			Runner: gardenRunner,
		},
		{
			Name:   "baggageclaim",
			Runner: baggageclaimRunner,
		},
	}

	if cmd.TSA.WorkerPrivateKey != "" {
		members = append(members, grouper.Member{
			Name:   "beacon",
			Runner: cmd.beaconRunner(logger.Session("beacon"), worker),
		})
	}

	runner := sigmon.New(grouper.NewParallel(os.Interrupt, members))

	return <-ifrit.Invoke(runner).Wait()
}

func (cmd *WorkerCommand) assetPath(paths ...string) string {
	return filepath.Join(append([]string{cmd.WorkDir.Path(), Version, "assets"}, paths...)...)
}

func (cmd *WorkerCommand) setup(logger lager.Logger) (bool, error) {
	okMarker := cmd.assetPath("ok")

	_, err := os.Stat(okMarker)
	if err == nil {
		logger.Info("already-done")
		return true, nil
	}

	logger.Info("unpacking")

	err = bindata.RestoreAssets(filepath.Split(cmd.assetPath()))
	if err != nil {
		logger.Error("failed-to-unpack", err)
		return false, err
	}

	_, err = os.Stat(cmd.assetPath())
	if os.IsNotExist(err) {
		logger.Info("no-assets")
		return false, nil
	}

	ok, err := os.Create(okMarker)
	if err != nil {
		logger.Error("failed-to-create-ok-marker", err)
		return false, err
	}

	err = ok.Close()
	if err != nil {
		logger.Error("failed-to-close-ok-marker", err)
		return false, err
	}

	logger.Info("done")

	return true, nil
}

func (cmd *WorkerCommand) workerName() (string, error) {
	if cmd.Name != "" {
		return cmd.Name, nil
	}

	return os.Hostname()
}

func (cmd *WorkerCommand) baggageclaimRunner(logger lager.Logger, hasAssets bool) (ifrit.Runner, error) {
	volumesDir := filepath.Join(cmd.WorkDir.Path(), "volumes")

	err := os.MkdirAll(volumesDir, 0755)
	if err != nil {
		return nil, err
	}

	cmd.Baggageclaim.Metrics = cmd.Metrics
	cmd.Baggageclaim.VolumesDir = baggageclaimcmd.DirFlag(volumesDir)

	cmd.Baggageclaim.OverlaysDir = filepath.Join(cmd.WorkDir.Path(), "overlays")

	if hasAssets {
		cmd.Baggageclaim.MkfsBin = cmd.assetPath("btrfs", "mkfs.btrfs")
		cmd.Baggageclaim.BtrfsBin = cmd.assetPath("btrfs", "btrfs")
	}

	return cmd.Baggageclaim.Runner(nil)
}

func (cmd *WorkerCommand) beaconRunner(logger lager.Logger, worker atc.Worker) ifrit.Runner {
	beacon := Beacon{
		Logger: logger,
		Config: cmd.TSA,
	}

	var beaconRunner ifrit.RunFunc
	if cmd.PeerIP != nil {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.PeerIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.PeerIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Register
	} else {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.BindIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.Baggageclaim.BindIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Forward
	}

	beacon.Worker = worker

	return restart.Restarter{
		Runner: beaconRunner,
		Load: func(prevRunner ifrit.Runner, prevErr error) ifrit.Runner {
			if prevErr == nil {
				return nil
			}

			if _, ok := prevErr.(*ssh.ExitError); !ok {
				logger.Error("restarting", prevErr)
				time.Sleep(5 * time.Second)
				return beaconRunner
			}

			return nil
		},
	}
}
<commit_msg>more realistic check for whether assets are there<commit_after>package main

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

	"golang.org/x/crypto/ssh"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	"github.com/concourse/bin/bindata"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/restart"
	"github.com/tedsuo/ifrit/sigmon"
)

type WorkerCommand struct {
	Name string   `long:"name" description:"The name to set for the worker during registration. If not specified, the hostname will be used."`
	Tags []string `long:"tag" description:"A tag to set during registration. Can be specified multiple times."`

	TeamName string `long:"team" description:"The name of the team that this worker will be assigned to."`

	HTTPProxy  URLFlag  `long:"http-proxy"  env:"http_proxy"                  description:"HTTP proxy endpoint to use for containers."`
	HTTPSProxy URLFlag  `long:"https-proxy" env:"https_proxy"                 description:"HTTPS proxy endpoint to use for containers."`
	NoProxy    []string `long:"no-proxy"    env:"no_proxy"    env-delim:","   description:"Blacklist of addresses to skip the proxy when reaching."`

	WorkDir DirFlag `long:"work-dir" required:"true" description:"Directory in which to place container data."`

	BindIP   IPFlag `long:"bind-ip"   default:"127.0.0.1" description:"IP address on which to listen for the Garden server."`
	BindPort uint16 `long:"bind-port" default:"7777"      description:"Port on which to listen for the Garden server."`

	PeerIP IPFlag `long:"peer-ip" description:"IP used to reach this worker from the ATC nodes. If omitted, the worker will be forwarded through the SSH connection to the TSA."`

	Garden GardenBackend `group:"Garden Configuration" namespace:"garden"`

	Baggageclaim baggageclaimcmd.BaggageclaimCommand `group:"Baggageclaim Configuration" namespace:"baggageclaim"`

	TSA BeaconConfig `group:"TSA Configuration" namespace:"tsa"`

	Metrics struct {
		YellerAPIKey      string `long:"yeller-api-key"     description:"Yeller API key. If specified, all errors logged will be emitted."`
		YellerEnvironment string `long:"yeller-environment" description:"Environment to tag on all Yeller events emitted."`
	} `group:"Metrics & Diagnostics"`
}

func (cmd *WorkerCommand) Execute(args []string) error {
	logger := lager.NewLogger("worker")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.INFO))

	hasAssets, err := cmd.setup(logger.Session("setup"))
	if err != nil {
		return err
	}

	worker, gardenRunner, err := cmd.gardenRunner(logger.Session("garden"), hasAssets)
	if err != nil {
		return err
	}

	worker.Version = WorkerVersion

	baggageclaimRunner, err := cmd.baggageclaimRunner(logger.Session("baggageclaim"), hasAssets)
	if err != nil {
		return err
	}

	members := grouper.Members{
		{
			Name:   "garden",
			Runner: gardenRunner,
		},
		{
			Name:   "baggageclaim",
			Runner: baggageclaimRunner,
		},
	}

	if cmd.TSA.WorkerPrivateKey != "" {
		members = append(members, grouper.Member{
			Name:   "beacon",
			Runner: cmd.beaconRunner(logger.Session("beacon"), worker),
		})
	}

	runner := sigmon.New(grouper.NewParallel(os.Interrupt, members))

	return <-ifrit.Invoke(runner).Wait()
}

func (cmd *WorkerCommand) assetPath(paths ...string) string {
	return filepath.Join(append([]string{cmd.WorkDir.Path(), Version, "assets"}, paths...)...)
}

func (cmd *WorkerCommand) setup(logger lager.Logger) (bool, error) {
	okMarker := cmd.assetPath("ok")

	_, err := os.Stat(okMarker)
	if err == nil {
		logger.Info("already-done")
		return true, nil
	}

	_, err = bindata.AssetDir("assets")
	if err != nil {
		logger.Info("no-assets")
		return false, nil
	}

	logger.Info("unpacking")

	err = bindata.RestoreAssets(filepath.Split(cmd.assetPath()))
	if err != nil {
		logger.Error("failed-to-unpack", err)
		return false, err
	}

	_, err = os.Stat(cmd.assetPath())
	if os.IsNotExist(err) {
		logger.Info("no-assets")
		return false, nil
	}

	ok, err := os.Create(okMarker)
	if err != nil {
		logger.Error("failed-to-create-ok-marker", err)
		return false, err
	}

	err = ok.Close()
	if err != nil {
		logger.Error("failed-to-close-ok-marker", err)
		return false, err
	}

	logger.Info("done")

	return true, nil
}

func (cmd *WorkerCommand) workerName() (string, error) {
	if cmd.Name != "" {
		return cmd.Name, nil
	}

	return os.Hostname()
}

func (cmd *WorkerCommand) baggageclaimRunner(logger lager.Logger, hasAssets bool) (ifrit.Runner, error) {
	volumesDir := filepath.Join(cmd.WorkDir.Path(), "volumes")

	err := os.MkdirAll(volumesDir, 0755)
	if err != nil {
		return nil, err
	}

	cmd.Baggageclaim.Metrics = cmd.Metrics
	cmd.Baggageclaim.VolumesDir = baggageclaimcmd.DirFlag(volumesDir)

	cmd.Baggageclaim.OverlaysDir = filepath.Join(cmd.WorkDir.Path(), "overlays")

	if hasAssets {
		cmd.Baggageclaim.MkfsBin = cmd.assetPath("btrfs", "mkfs.btrfs")
		cmd.Baggageclaim.BtrfsBin = cmd.assetPath("btrfs", "btrfs")
	}

	return cmd.Baggageclaim.Runner(nil)
}

func (cmd *WorkerCommand) beaconRunner(logger lager.Logger, worker atc.Worker) ifrit.Runner {
	beacon := Beacon{
		Logger: logger,
		Config: cmd.TSA,
	}

	var beaconRunner ifrit.RunFunc
	if cmd.PeerIP != nil {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.PeerIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.PeerIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Register
	} else {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.BindIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.Baggageclaim.BindIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Forward
	}

	beacon.Worker = worker

	return restart.Restarter{
		Runner: beaconRunner,
		Load: func(prevRunner ifrit.Runner, prevErr error) ifrit.Runner {
			if prevErr == nil {
				return nil
			}

			if _, ok := prevErr.(*ssh.ExitError); !ok {
				logger.Error("restarting", prevErr)
				time.Sleep(5 * time.Second)
				return beaconRunner
			}

			return nil
		},
	}
}
<|endoftext|>