<commit_before>// Copyright 2015 The Serviced 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 elasticsearch

import (
	"fmt"
	"os"

	"path/filepath"
	"sync"
	"time"

	model "github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/dfs"
	"github.com/control-center/serviced/volume"
	gzip "github.com/klauspost/pgzip"
	"github.com/zenoss/glog"
)

// InProgress prompts which backup is currently backing up or restoring
type InProgress struct {
	locker   *sync.RWMutex
	running  bool
	filename string
	op       string
	err      error
}

// SetProgress sets the current operation
func (p *InProgress) SetProgress(filename, op string) {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.running = true
	p.filename = filename
	p.op = op
}

// SetError sets the error returned from the current operation
func (p *InProgress) SetError(err error) {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.running = false
	p.err = err
}

// Reset resets the progress indicator
func (p *InProgress) Reset() {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.running = false
	p.filename = ""
	p.op = ""
	p.err = nil
}

// UnsetProgress unsets the current operation
func (p *InProgress) UnsetProgress(err error) {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.filename = ""
	p.op = ""
	p.err = err
}

// GetProgress returns the progress of the current running backup or restore.
func (p *InProgress) GetProgress() (bool, string, string, error) {
	p.locker.RLock()
	defer p.locker.RUnlock()
	return p.running, p.filename, p.op, p.err
}

var inprogress = &InProgress{locker: &sync.RWMutex{}}

// Backup takes a backup of the full application stack and returns the filename
// that it is written to.
func (dao *ControlPlaneDao) Backup(backupRequest model.BackupRequest, filename *string) (err error) {
	ctx := datastore.Get()

	dirpath := backupRequest.Dirpath
	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("backup")
	defer dfslocker.Unlock()

	// set the progress of the backup file
	*filename = time.Now().UTC().Format("backup-2006-01-02-150405.tgz")
	backupfilename := filepath.Join(dirpath, *filename)
	if dirpath == "" {
		backupfilename = filepath.Join(dao.backupsPath, *filename)
	}
	inprogress.SetProgress(backupfilename, "backup")
	defer func() {
		if err != nil {
			glog.Errorf("Backup failed with error: %s", err)
			os.Remove(backupfilename)
		}
		inprogress.SetError(err)
	}()
	// create the file and write
	fh, err := os.Create(backupfilename)
	if err != nil {
		glog.Errorf("Could not create backup file at %s: %s", backupfilename, err)
		return
	}
	defer fh.Close()
	w := gzip.NewWriter(fh)
	defer w.Close()
	err = dao.facade.Backup(ctx, w, backupRequest.SnapshotSpacePercent)
	return
}

// AsyncBackup is the same as backup, but asynchronous
func (dao *ControlPlaneDao) AsyncBackup(backupRequest model.BackupRequest, filename *string) (err error) {
	ctx := datastore.Get()
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("backup")
	inprogress.Reset()
	dfslocker.Unlock()
	go dao.Backup(backupRequest, filename)
	return
}

// Restore restores the full application stack from a backup file.
func (dao *ControlPlaneDao) Restore(filename string, _ *int) (err error) {
	ctx := datastore.Get()

	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("restore")
	defer dfslocker.Unlock()
	inprogress.SetProgress(filename, "restore")
	defer func() {
		if err != nil {
			glog.Errorf("Restore failed with error: %s", err)
		}
		inprogress.SetError(err)
	}()
	info, err := dfs.ExtractBackupInfo(filename)
	if err != nil {
		return err
	}
	fh, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fh.Close()
	gz, err := gzip.NewReader(fh)
	if err != nil {
		return err
	}
	defer gz.Close()
	err = dao.facade.Restore(ctx, gz, info)
	return err
}

// AsyncRestore is the same as restore, but asynchronous.
func (dao *ControlPlaneDao) AsyncRestore(filename string, unused *int) (err error) {
	ctx := datastore.Get()
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("restore")
	inprogress.Reset()
	dfslocker.Unlock()
	go dao.Restore(filename, unused)
	return
}

// ListBackups returns the list of backups
func (dao *ControlPlaneDao) ListBackups(dirpath string, files *[]model.BackupFile) (err error) {
	*files = []model.BackupFile{}

	// Read the contents of the directory
	if dirpath == "" {
		dirpath = dao.backupsPath
	}
	dir, err := os.Open(dirpath)
	if err != nil {
		return err
	}
	defer dir.Close()

	fis, err := dir.Readdir(0)
	if err != nil {
		return err
	}
	// What is currently running?
	running, fp, _, _ := inprogress.GetProgress()

	for _, fi := range fis {
		if !fi.IsDir() {
			// Set up the backup file
			fullpath := filepath.Join(dirpath, fi.Name())
			bf := model.BackupFile{
				InProgress: running && fullpath == fp,
				FullPath:   fullpath,
				Name:       fi.Name(),
				Size:       fi.Size(),
				Mode:       fi.Mode(),
				ModTime:    fi.ModTime(),
			}
			// If it is not running, make sure the backup is legit
			if !bf.InProgress {
				isbackup := func(filename string) bool {
					fh, err := os.Open(fullpath)
					if err != nil {
						return false
					}
					defer fh.Close()
					gz, err := gzip.NewReader(fh)
					if err != nil {
						return false
					}
					defer gz.Close()
					_, err = dao.facade.BackupInfo(datastore.Get(), gz)
					if err != nil {
						return false
					}
					return true
				}(fullpath)
				if !isbackup {
					continue
				}
			}
			// Add the file to the list
			*files = append(*files, bf)
		}
	}
	return
}

// BackupStatus returns the current status of the backup or restore that is
// running.
func (dao *ControlPlaneDao) BackupStatus(_ model.EntityRequest, status *string) (err error) {
	running, filename, op, err := inprogress.GetProgress()
	if running {
		*status = fmt.Sprintf("Performing a %s on %s", op, filename)
	} else {
		if err != nil {
			*status = fmt.Sprintf("Completed a %s on %s with error: %s", op, filename, err)
		} else {
			inprogress.Reset()
			*status = ""
		}
	}
	return
}

// Snapshot captures the current state of a single application
func (dao *ControlPlaneDao) Snapshot(req model.SnapshotRequest, snapshotID *string) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("snapshot")
	defer dfslocker.Unlock()

	tagList := []string{}
	if len(req.Tag) > 0 {
		tagList = []string{req.Tag}
	}

	if req.ContainerID != "" {
		*snapshotID, err = dao.facade.Commit(ctx, req.ContainerID, req.Message, tagList, req.SnapshotSpacePercent)
	} else {
		*snapshotID, err = dao.facade.Snapshot(ctx, req.ServiceID, req.Message, tagList, req.SnapshotSpacePercent)
	}
	return
}

// Rollback reverts a single application to a particular state
func (dao *ControlPlaneDao) Rollback(req model.RollbackRequest, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("rollback")
	defer dfslocker.Unlock()
	err = dao.facade.Rollback(ctx, req.SnapshotID, req.ForceRestart)
	return
}

// DeleteSnapshot deletes a single snapshot
func (dao *ControlPlaneDao) DeleteSnapshot(snapshotID string, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("delete snapshot")
	defer dfslocker.Unlock()
	err = dao.facade.DeleteSnapshot(ctx, snapshotID)
	return
}

// DeleteSnapshots deletes all snapshots for a service
func (dao *ControlPlaneDao) DeleteSnapshots(serviceID string, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("delete snapshots")
	defer dfslocker.Unlock()
	err = dao.facade.DeleteSnapshots(ctx, serviceID)
	return
}

// ListSnapshots returns a list of all snapshots for a service
func (dao *ControlPlaneDao) ListSnapshots(serviceID string, snapshots *[]model.SnapshotInfo) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("list snapshots")
	defer dfslocker.Unlock()

	*snapshots = make([]model.SnapshotInfo, 0)
	snapshotIDs, err := dao.facade.ListSnapshots(ctx, serviceID)
	if err != nil {
		return err
	}
	for _, snapshotID := range snapshotIDs {
		var newInfo model.SnapshotInfo

		info, err := dao.facade.GetSnapshotInfo(ctx, snapshotID)
		if err == volume.ErrInvalidSnapshot {
			newInfo = model.SnapshotInfo{
				SnapshotID: snapshotID,
				Invalid:    true,
			}

		} else if err != nil {
			return err
		} else {
			newInfo = model.SnapshotInfo{
				SnapshotID:  info.Name,
				TenantID:    info.TenantID,
				Description: info.Message,
				Tags:        info.Tags,
				Created:     info.Created,
				Invalid:     false,
			}
		}
		*snapshots = append(*snapshots, newInfo)
	}
	return
}

// ResetRegistry prompts all images to be pushed back into the docker registry
func (dao *ControlPlaneDao) ResetRegistry(_ model.EntityRequest, _ *int) (err error) {
	// Do not DFSLock here, Facade does that
	err = dao.facade.SyncRegistryImages(datastore.Get(), true)
	return
}

// RepairRegistry will try to recover the latest image of all service images
// from the docker registry and save it to the index.
func (dao *ControlPlaneDao) RepairRegistry(_ model.EntityRequest, _ *int) (err error) {
	// Do not DFSLock here, Facade does that
	err = dao.facade.RepairRegistry(datastore.Get())
	return
}

// ReadyDFS locks until it receives notice that the dfs is idle
func (dao *ControlPlaneDao) ReadyDFS(serviceID string, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("reset tenant lock")
	defer dfslocker.Unlock()

	err = dao.facade.ResetLock(ctx, serviceID)
	return
}

// TagSnapshot adds a tag to an existing snapshot
func (dao *ControlPlaneDao) TagSnapshot(request model.TagSnapshotRequest, _ *int) error {
	return dao.facade.TagSnapshot(request.SnapshotID, request.TagName)
}

// RemoveSnapshotTag removes a tag from an existing snapshot
func (dao *ControlPlaneDao) RemoveSnapshotTag(request model.SnapshotByTagRequest, snapshotID *string) (err error) {
	ctx := datastore.Get()
	*snapshotID, err = dao.facade.RemoveSnapshotTag(ctx, request.ServiceID, request.TagName)
	return err
}

// GetSnapshotByServiceIDAndTag Gets the snapshot from a specific service with a specific tag
func (dao *ControlPlaneDao) GetSnapshotByServiceIDAndTag(request model.SnapshotByTagRequest, snapshot *model.SnapshotInfo) (err error) {
	ctx := datastore.Get()
	info, err := dao.facade.GetSnapshotByServiceIDAndTag(ctx, request.ServiceID, request.TagName)
	if err != nil {
		return
	}
	*snapshot = model.SnapshotInfo{
		SnapshotID:  info.Name,
		TenantID:    info.TenantID,
		Description: info.Message,
		Tags:        info.Tags,
		Created:     info.Created,
	}
	return nil
}
<commit_msg>CC-2292: Set backup concurrency to 2<commit_after>// Copyright 2015 The Serviced 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 elasticsearch

import (
	"fmt"
	"os"

	"path/filepath"
	"sync"
	"time"

	model "github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/dfs"
	"github.com/control-center/serviced/volume"
	gzip "github.com/klauspost/pgzip"
	"github.com/zenoss/glog"
)

// InProgress prompts which backup is currently backing up or restoring
type InProgress struct {
	locker   *sync.RWMutex
	running  bool
	filename string
	op       string
	err      error
}

// SetProgress sets the current operation
func (p *InProgress) SetProgress(filename, op string) {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.running = true
	p.filename = filename
	p.op = op
}

// SetError sets the error returned from the current operation
func (p *InProgress) SetError(err error) {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.running = false
	p.err = err
}

// Reset resets the progress indicator
func (p *InProgress) Reset() {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.running = false
	p.filename = ""
	p.op = ""
	p.err = nil
}

// UnsetProgress unsets the current operation
func (p *InProgress) UnsetProgress(err error) {
	p.locker.Lock()
	defer p.locker.Unlock()
	p.filename = ""
	p.op = ""
	p.err = err
}

// GetProgress returns the progress of the current running backup or restore.
func (p *InProgress) GetProgress() (bool, string, string, error) {
	p.locker.RLock()
	defer p.locker.RUnlock()
	return p.running, p.filename, p.op, p.err
}

var inprogress = &InProgress{locker: &sync.RWMutex{}}

// Backup takes a backup of the full application stack and returns the filename
// that it is written to.
func (dao *ControlPlaneDao) Backup(backupRequest model.BackupRequest, filename *string) (err error) {
	ctx := datastore.Get()

	dirpath := backupRequest.Dirpath
	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("backup")
	defer dfslocker.Unlock()

	// set the progress of the backup file
	*filename = time.Now().UTC().Format("backup-2006-01-02-150405.tgz")
	backupfilename := filepath.Join(dirpath, *filename)
	if dirpath == "" {
		backupfilename = filepath.Join(dao.backupsPath, *filename)
	}
	inprogress.SetProgress(backupfilename, "backup")
	defer func() {
		if err != nil {
			glog.Errorf("Backup failed with error: %s", err)
			os.Remove(backupfilename)
		}
		inprogress.SetError(err)
	}()
	// create the file and write
	fh, err := os.Create(backupfilename)
	if err != nil {
		glog.Errorf("Could not create backup file at %s: %s", backupfilename, err)
		return
	}
	defer fh.Close()
	w := gzip.NewWriter(fh)
	// CC-2292: Limit concurrency of backup gzipping
	// This setting will cause the writer to process up to 2 100KB blocks
	// at a time before the writer blocks. The default was 16 250KB blocks.
	// Smaller blocks will allow other goroutines to get time more frequently.
	w.SetConcurrency(100000, 2)
	defer w.Close()
	err = dao.facade.Backup(ctx, w, backupRequest.SnapshotSpacePercent)
	return
}

// AsyncBackup is the same as backup, but asynchronous
func (dao *ControlPlaneDao) AsyncBackup(backupRequest model.BackupRequest, filename *string) (err error) {
	ctx := datastore.Get()
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("backup")
	inprogress.Reset()
	dfslocker.Unlock()
	go dao.Backup(backupRequest, filename)
	return
}

// Restore restores the full application stack from a backup file.
func (dao *ControlPlaneDao) Restore(filename string, _ *int) (err error) {
	ctx := datastore.Get()

	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("restore")
	defer dfslocker.Unlock()
	inprogress.SetProgress(filename, "restore")
	defer func() {
		if err != nil {
			glog.Errorf("Restore failed with error: %s", err)
		}
		inprogress.SetError(err)
	}()
	info, err := dfs.ExtractBackupInfo(filename)
	if err != nil {
		return err
	}
	fh, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fh.Close()
	gz, err := gzip.NewReader(fh)
	if err != nil {
		return err
	}
	defer gz.Close()
	err = dao.facade.Restore(ctx, gz, info)
	return err
}

// AsyncRestore is the same as restore, but asynchronous.
func (dao *ControlPlaneDao) AsyncRestore(filename string, unused *int) (err error) {
	ctx := datastore.Get()
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("restore")
	inprogress.Reset()
	dfslocker.Unlock()
	go dao.Restore(filename, unused)
	return
}

// ListBackups returns the list of backups
func (dao *ControlPlaneDao) ListBackups(dirpath string, files *[]model.BackupFile) (err error) {
	*files = []model.BackupFile{}

	// Read the contents of the directory
	if dirpath == "" {
		dirpath = dao.backupsPath
	}
	dir, err := os.Open(dirpath)
	if err != nil {
		return err
	}
	defer dir.Close()

	fis, err := dir.Readdir(0)
	if err != nil {
		return err
	}
	// What is currently running?
	running, fp, _, _ := inprogress.GetProgress()

	for _, fi := range fis {
		if !fi.IsDir() {
			// Set up the backup file
			fullpath := filepath.Join(dirpath, fi.Name())
			bf := model.BackupFile{
				InProgress: running && fullpath == fp,
				FullPath:   fullpath,
				Name:       fi.Name(),
				Size:       fi.Size(),
				Mode:       fi.Mode(),
				ModTime:    fi.ModTime(),
			}
			// If it is not running, make sure the backup is legit
			if !bf.InProgress {
				isbackup := func(filename string) bool {
					fh, err := os.Open(fullpath)
					if err != nil {
						return false
					}
					defer fh.Close()
					gz, err := gzip.NewReader(fh)
					if err != nil {
						return false
					}
					defer gz.Close()
					_, err = dao.facade.BackupInfo(datastore.Get(), gz)
					if err != nil {
						return false
					}
					return true
				}(fullpath)
				if !isbackup {
					continue
				}
			}
			// Add the file to the list
			*files = append(*files, bf)
		}
	}
	return
}

// BackupStatus returns the current status of the backup or restore that is
// running.
func (dao *ControlPlaneDao) BackupStatus(_ model.EntityRequest, status *string) (err error) {
	running, filename, op, err := inprogress.GetProgress()
	if running {
		*status = fmt.Sprintf("Performing a %s on %s", op, filename)
	} else {
		if err != nil {
			*status = fmt.Sprintf("Completed a %s on %s with error: %s", op, filename, err)
		} else {
			inprogress.Reset()
			*status = ""
		}
	}
	return
}

// Snapshot captures the current state of a single application
func (dao *ControlPlaneDao) Snapshot(req model.SnapshotRequest, snapshotID *string) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("snapshot")
	defer dfslocker.Unlock()

	tagList := []string{}
	if len(req.Tag) > 0 {
		tagList = []string{req.Tag}
	}

	if req.ContainerID != "" {
		*snapshotID, err = dao.facade.Commit(ctx, req.ContainerID, req.Message, tagList, req.SnapshotSpacePercent)
	} else {
		*snapshotID, err = dao.facade.Snapshot(ctx, req.ServiceID, req.Message, tagList, req.SnapshotSpacePercent)
	}
	return
}

// Rollback reverts a single application to a particular state
func (dao *ControlPlaneDao) Rollback(req model.RollbackRequest, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("rollback")
	defer dfslocker.Unlock()
	err = dao.facade.Rollback(ctx, req.SnapshotID, req.ForceRestart)
	return
}

// DeleteSnapshot deletes a single snapshot
func (dao *ControlPlaneDao) DeleteSnapshot(snapshotID string, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("delete snapshot")
	defer dfslocker.Unlock()
	err = dao.facade.DeleteSnapshot(ctx, snapshotID)
	return
}

// DeleteSnapshots deletes all snapshots for a service
func (dao *ControlPlaneDao) DeleteSnapshots(serviceID string, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("delete snapshots")
	defer dfslocker.Unlock()
	err = dao.facade.DeleteSnapshots(ctx, serviceID)
	return
}

// ListSnapshots returns a list of all snapshots for a service
func (dao *ControlPlaneDao) ListSnapshots(serviceID string, snapshots *[]model.SnapshotInfo) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("list snapshots")
	defer dfslocker.Unlock()

	*snapshots = make([]model.SnapshotInfo, 0)
	snapshotIDs, err := dao.facade.ListSnapshots(ctx, serviceID)
	if err != nil {
		return err
	}
	for _, snapshotID := range snapshotIDs {
		var newInfo model.SnapshotInfo

		info, err := dao.facade.GetSnapshotInfo(ctx, snapshotID)
		if err == volume.ErrInvalidSnapshot {
			newInfo = model.SnapshotInfo{
				SnapshotID: snapshotID,
				Invalid:    true,
			}

		} else if err != nil {
			return err
		} else {
			newInfo = model.SnapshotInfo{
				SnapshotID:  info.Name,
				TenantID:    info.TenantID,
				Description: info.Message,
				Tags:        info.Tags,
				Created:     info.Created,
				Invalid:     false,
			}
		}
		*snapshots = append(*snapshots, newInfo)
	}
	return
}

// ResetRegistry prompts all images to be pushed back into the docker registry
func (dao *ControlPlaneDao) ResetRegistry(_ model.EntityRequest, _ *int) (err error) {
	// Do not DFSLock here, Facade does that
	err = dao.facade.SyncRegistryImages(datastore.Get(), true)
	return
}

// RepairRegistry will try to recover the latest image of all service images
// from the docker registry and save it to the index.
func (dao *ControlPlaneDao) RepairRegistry(_ model.EntityRequest, _ *int) (err error) {
	// Do not DFSLock here, Facade does that
	err = dao.facade.RepairRegistry(datastore.Get())
	return
}

// ReadyDFS locks until it receives notice that the dfs is idle
func (dao *ControlPlaneDao) ReadyDFS(serviceID string, _ *int) (err error) {
	ctx := datastore.Get()

	// synchronize the dfs
	dfslocker := dao.facade.DFSLock(ctx)
	dfslocker.Lock("reset tenant lock")
	defer dfslocker.Unlock()

	err = dao.facade.ResetLock(ctx, serviceID)
	return
}

// TagSnapshot adds a tag to an existing snapshot
func (dao *ControlPlaneDao) TagSnapshot(request model.TagSnapshotRequest, _ *int) error {
	return dao.facade.TagSnapshot(request.SnapshotID, request.TagName)
}

// RemoveSnapshotTag removes a tag from an existing snapshot
func (dao *ControlPlaneDao) RemoveSnapshotTag(request model.SnapshotByTagRequest, snapshotID *string) (err error) {
	ctx := datastore.Get()
	*snapshotID, err = dao.facade.RemoveSnapshotTag(ctx, request.ServiceID, request.TagName)
	return err
}

// GetSnapshotByServiceIDAndTag Gets the snapshot from a specific service with a specific tag
func (dao *ControlPlaneDao) GetSnapshotByServiceIDAndTag(request model.SnapshotByTagRequest, snapshot *model.SnapshotInfo) (err error) {
	ctx := datastore.Get()
	info, err := dao.facade.GetSnapshotByServiceIDAndTag(ctx, request.ServiceID, request.TagName)
	if err != nil {
		return
	}
	*snapshot = model.SnapshotInfo{
		SnapshotID:  info.Name,
		TenantID:    info.TenantID,
		Description: info.Message,
		Tags:        info.Tags,
		Created:     info.Created,
	}
	return nil
}
<|endoftext|>


<commit_before>package quic

import (
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/flowcontrol"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
)

type receiveStreamI interface {
	ReceiveStream

	handleStreamFrame(*wire.StreamFrame) error
	handleResetStreamFrame(*wire.ResetStreamFrame) error
	closeForShutdown(error)
	getWindowUpdate() protocol.ByteCount
}

type receiveStream struct {
	mutex sync.Mutex

	streamID protocol.StreamID

	sender streamSender

	frameQueue  *frameSorter
	readOffset  protocol.ByteCount
	finalOffset protocol.ByteCount

	currentFrame       []byte
	currentFrameIsLast bool // is the currentFrame the last frame on this stream
	readPosInFrame     int

	closeForShutdownErr error
	cancelReadErr       error
	resetRemotelyErr    StreamError

	closedForShutdown bool // set when CloseForShutdown() is called
	finRead           bool // set once we read a frame with a FinBit
	canceledRead      bool // set when CancelRead() is called
	resetRemotely     bool // set when HandleResetStreamFrame() is called

	readChan chan struct{}
	deadline time.Time

	flowController flowcontrol.StreamFlowController
	version        protocol.VersionNumber
}

var _ ReceiveStream = &receiveStream{}
var _ receiveStreamI = &receiveStream{}

func newReceiveStream(
	streamID protocol.StreamID,
	sender streamSender,
	flowController flowcontrol.StreamFlowController,
	version protocol.VersionNumber,
) *receiveStream {
	return &receiveStream{
		streamID:       streamID,
		sender:         sender,
		flowController: flowController,
		frameQueue:     newFrameSorter(),
		readChan:       make(chan struct{}, 1),
		finalOffset:    protocol.MaxByteCount,
		version:        version,
	}
}

func (s *receiveStream) StreamID() protocol.StreamID {
	return s.streamID
}

// Read implements io.Reader. It is not thread safe!
func (s *receiveStream) Read(p []byte) (int, error) {
	completed, n, err := s.readImpl(p)
	if completed {
		s.streamCompleted()
	}
	return n, err
}

func (s *receiveStream) readImpl(p []byte) (bool /*stream completed */, int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.finRead {
		return false, 0, io.EOF
	}
	if s.canceledRead {
		return false, 0, s.cancelReadErr
	}
	if s.resetRemotely {
		return false, 0, s.resetRemotelyErr
	}
	if s.closedForShutdown {
		return false, 0, s.closeForShutdownErr
	}

	bytesRead := 0
	for bytesRead < len(p) {
		if s.currentFrame == nil || s.readPosInFrame >= len(s.currentFrame) {
			s.dequeueNextFrame()
		}
		if s.currentFrame == nil && bytesRead > 0 {
			return false, bytesRead, s.closeForShutdownErr
		}

		var deadlineTimer *utils.Timer
		for {
			// Stop waiting on errors
			if s.closedForShutdown {
				return false, bytesRead, s.closeForShutdownErr
			}
			if s.canceledRead {
				return false, bytesRead, s.cancelReadErr
			}
			if s.resetRemotely {
				return false, bytesRead, s.resetRemotelyErr
			}

			deadline := s.deadline
			if !deadline.IsZero() {
				if !time.Now().Before(deadline) {
					return false, bytesRead, errDeadline
				}
				if deadlineTimer == nil {
					deadlineTimer = utils.NewTimer()
				}
				deadlineTimer.Reset(deadline)
			}

			if s.currentFrame != nil || s.currentFrameIsLast {
				break
			}

			s.mutex.Unlock()
			if deadline.IsZero() {
				<-s.readChan
			} else {
				select {
				case <-s.readChan:
				case <-deadlineTimer.Chan():
					deadlineTimer.SetRead()
				}
			}
			s.mutex.Lock()
			if s.currentFrame == nil {
				s.dequeueNextFrame()
			}
		}

		if bytesRead > len(p) {
			return false, bytesRead, fmt.Errorf("BUG: bytesRead (%d) > len(p) (%d) in stream.Read", bytesRead, len(p))
		}
		if s.readPosInFrame > len(s.currentFrame) {
			return false, bytesRead, fmt.Errorf("BUG: readPosInFrame (%d) > frame.DataLen (%d) in stream.Read", s.readPosInFrame, len(s.currentFrame))
		}

		s.mutex.Unlock()

		m := copy(p[bytesRead:], s.currentFrame[s.readPosInFrame:])
		s.readPosInFrame += m
		bytesRead += m
		s.readOffset += protocol.ByteCount(m)

		s.mutex.Lock()
		// when a RESET_STREAM was received, the was already informed about the final byteOffset for this stream
		if !s.resetRemotely {
			s.flowController.AddBytesRead(protocol.ByteCount(m))
		}

		if s.readPosInFrame >= len(s.currentFrame) && s.currentFrameIsLast {
			s.finRead = true
			return true, bytesRead, io.EOF
		}
	}
	return false, bytesRead, nil
}

func (s *receiveStream) dequeueNextFrame() {
	var offset protocol.ByteCount
	offset, s.currentFrame = s.frameQueue.Pop()
	s.currentFrameIsLast = offset+protocol.ByteCount(len(s.currentFrame)) >= s.finalOffset
	s.readPosInFrame = 0
}

func (s *receiveStream) CancelRead(errorCode protocol.ApplicationErrorCode) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.finRead || s.canceledRead || s.resetRemotely {
		return
	}
	if s.finalOffset != protocol.MaxByteCount { // final offset was already received
		s.streamCompleted()
	}
	s.canceledRead = true
	s.cancelReadErr = fmt.Errorf("Read on stream %d canceled with error code %d", s.streamID, errorCode)
	s.signalRead()
	s.sender.queueControlFrame(&wire.StopSendingFrame{
		StreamID:  s.streamID,
		ErrorCode: errorCode,
	})
}

func (s *receiveStream) handleStreamFrame(frame *wire.StreamFrame) error {
	maxOffset := frame.Offset + frame.DataLen()

	s.mutex.Lock()
	defer s.mutex.Unlock()

	if err := s.flowController.UpdateHighestReceived(maxOffset, frame.FinBit); err != nil {
		return err
	}
	if frame.FinBit {
		s.finalOffset = maxOffset
	}
	if s.canceledRead {
		if frame.FinBit {
			s.streamCompleted()
		}
		return nil
	}
	if err := s.frameQueue.Push(frame.Data, frame.Offset); err != nil {
		return err
	}
	if frame.FinBit {
		s.finalOffset = maxOffset
	}
	s.signalRead()
	return nil
}

func (s *receiveStream) handleResetStreamFrame(frame *wire.ResetStreamFrame) error {
	completed, err := s.handleResetStreamFrameImpl(frame)
	if completed {
		s.streamCompleted()
	}
	return err
}

func (s *receiveStream) handleResetStreamFrameImpl(frame *wire.ResetStreamFrame) (bool /*completed */, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.closedForShutdown {
		return false, nil
	}
	if err := s.flowController.UpdateHighestReceived(frame.ByteOffset, true); err != nil {
		return false, err
	}
	s.finalOffset = frame.ByteOffset

	// ignore duplicate RESET_STREAM frames for this stream (after checking their final offset)
	if s.resetRemotely {
		return false, nil
	}
	s.resetRemotely = true
	s.resetRemotelyErr = streamCanceledError{
		errorCode: frame.ErrorCode,
		error:     fmt.Errorf("Stream %d was reset with error code %d", s.streamID, frame.ErrorCode),
	}
	s.signalRead()
	return true, nil
}

func (s *receiveStream) CloseRemote(offset protocol.ByteCount) {
	s.handleStreamFrame(&wire.StreamFrame{FinBit: true, Offset: offset})
}

func (s *receiveStream) SetReadDeadline(t time.Time) error {
	s.mutex.Lock()
	s.deadline = t
	s.mutex.Unlock()
	s.signalRead()
	return nil
}

// CloseForShutdown closes a stream abruptly.
// It makes Read unblock (and return the error) immediately.
// The peer will NOT be informed about this: the stream is closed without sending a FIN or RESET.
func (s *receiveStream) closeForShutdown(err error) {
	s.mutex.Lock()
	s.closedForShutdown = true
	s.closeForShutdownErr = err
	s.mutex.Unlock()
	s.signalRead()
}

func (s *receiveStream) getWindowUpdate() protocol.ByteCount {
	return s.flowController.GetWindowUpdate()
}

func (s *receiveStream) streamCompleted() {
	if !s.finRead {
		s.flowController.Abandon()
	}
	s.sender.onStreamCompleted(s.streamID)
}

// signalRead performs a non-blocking send on the readChan
func (s *receiveStream) signalRead() {
	select {
	case s.readChan <- struct{}{}:
	default:
	}
}
<commit_msg>call the onStreamCompleted callback without holding the stream mutex<commit_after>package quic

import (
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/flowcontrol"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
)

type receiveStreamI interface {
	ReceiveStream

	handleStreamFrame(*wire.StreamFrame) error
	handleResetStreamFrame(*wire.ResetStreamFrame) error
	closeForShutdown(error)
	getWindowUpdate() protocol.ByteCount
}

type receiveStream struct {
	mutex sync.Mutex

	streamID protocol.StreamID

	sender streamSender

	frameQueue  *frameSorter
	readOffset  protocol.ByteCount
	finalOffset protocol.ByteCount

	currentFrame       []byte
	currentFrameIsLast bool // is the currentFrame the last frame on this stream
	readPosInFrame     int

	closeForShutdownErr error
	cancelReadErr       error
	resetRemotelyErr    StreamError

	closedForShutdown bool // set when CloseForShutdown() is called
	finRead           bool // set once we read a frame with a FinBit
	canceledRead      bool // set when CancelRead() is called
	resetRemotely     bool // set when HandleResetStreamFrame() is called

	readChan chan struct{}
	deadline time.Time

	flowController flowcontrol.StreamFlowController
	version        protocol.VersionNumber
}

var _ ReceiveStream = &receiveStream{}
var _ receiveStreamI = &receiveStream{}

func newReceiveStream(
	streamID protocol.StreamID,
	sender streamSender,
	flowController flowcontrol.StreamFlowController,
	version protocol.VersionNumber,
) *receiveStream {
	return &receiveStream{
		streamID:       streamID,
		sender:         sender,
		flowController: flowController,
		frameQueue:     newFrameSorter(),
		readChan:       make(chan struct{}, 1),
		finalOffset:    protocol.MaxByteCount,
		version:        version,
	}
}

func (s *receiveStream) StreamID() protocol.StreamID {
	return s.streamID
}

// Read implements io.Reader. It is not thread safe!
func (s *receiveStream) Read(p []byte) (int, error) {
	completed, n, err := s.readImpl(p)
	if completed {
		s.streamCompleted()
	}
	return n, err
}

func (s *receiveStream) readImpl(p []byte) (bool /*stream completed */, int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.finRead {
		return false, 0, io.EOF
	}
	if s.canceledRead {
		return false, 0, s.cancelReadErr
	}
	if s.resetRemotely {
		return false, 0, s.resetRemotelyErr
	}
	if s.closedForShutdown {
		return false, 0, s.closeForShutdownErr
	}

	bytesRead := 0
	for bytesRead < len(p) {
		if s.currentFrame == nil || s.readPosInFrame >= len(s.currentFrame) {
			s.dequeueNextFrame()
		}
		if s.currentFrame == nil && bytesRead > 0 {
			return false, bytesRead, s.closeForShutdownErr
		}

		var deadlineTimer *utils.Timer
		for {
			// Stop waiting on errors
			if s.closedForShutdown {
				return false, bytesRead, s.closeForShutdownErr
			}
			if s.canceledRead {
				return false, bytesRead, s.cancelReadErr
			}
			if s.resetRemotely {
				return false, bytesRead, s.resetRemotelyErr
			}

			deadline := s.deadline
			if !deadline.IsZero() {
				if !time.Now().Before(deadline) {
					return false, bytesRead, errDeadline
				}
				if deadlineTimer == nil {
					deadlineTimer = utils.NewTimer()
				}
				deadlineTimer.Reset(deadline)
			}

			if s.currentFrame != nil || s.currentFrameIsLast {
				break
			}

			s.mutex.Unlock()
			if deadline.IsZero() {
				<-s.readChan
			} else {
				select {
				case <-s.readChan:
				case <-deadlineTimer.Chan():
					deadlineTimer.SetRead()
				}
			}
			s.mutex.Lock()
			if s.currentFrame == nil {
				s.dequeueNextFrame()
			}
		}

		if bytesRead > len(p) {
			return false, bytesRead, fmt.Errorf("BUG: bytesRead (%d) > len(p) (%d) in stream.Read", bytesRead, len(p))
		}
		if s.readPosInFrame > len(s.currentFrame) {
			return false, bytesRead, fmt.Errorf("BUG: readPosInFrame (%d) > frame.DataLen (%d) in stream.Read", s.readPosInFrame, len(s.currentFrame))
		}

		s.mutex.Unlock()

		m := copy(p[bytesRead:], s.currentFrame[s.readPosInFrame:])
		s.readPosInFrame += m
		bytesRead += m
		s.readOffset += protocol.ByteCount(m)

		s.mutex.Lock()
		// when a RESET_STREAM was received, the was already informed about the final byteOffset for this stream
		if !s.resetRemotely {
			s.flowController.AddBytesRead(protocol.ByteCount(m))
		}

		if s.readPosInFrame >= len(s.currentFrame) && s.currentFrameIsLast {
			s.finRead = true
			return true, bytesRead, io.EOF
		}
	}
	return false, bytesRead, nil
}

func (s *receiveStream) dequeueNextFrame() {
	var offset protocol.ByteCount
	offset, s.currentFrame = s.frameQueue.Pop()
	s.currentFrameIsLast = offset+protocol.ByteCount(len(s.currentFrame)) >= s.finalOffset
	s.readPosInFrame = 0
}

func (s *receiveStream) CancelRead(errorCode protocol.ApplicationErrorCode) {
	if completed := s.cancelReadImpl(errorCode); completed {
		s.streamCompleted()
	}
}

func (s *receiveStream) cancelReadImpl(errorCode protocol.ApplicationErrorCode) bool /* completed */ {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.finRead || s.canceledRead || s.resetRemotely {
		return false
	}
	s.canceledRead = true
	s.cancelReadErr = fmt.Errorf("Read on stream %d canceled with error code %d", s.streamID, errorCode)
	s.signalRead()
	s.sender.queueControlFrame(&wire.StopSendingFrame{
		StreamID:  s.streamID,
		ErrorCode: errorCode,
	})
	// We're done with this stream if the final offset was already received.
	return s.finalOffset != protocol.MaxByteCount
}

func (s *receiveStream) handleStreamFrame(frame *wire.StreamFrame) error {
	completed, err := s.handleStreamFrameImpl(frame)
	if completed {
		s.streamCompleted()
	}
	return err
}

func (s *receiveStream) handleStreamFrameImpl(frame *wire.StreamFrame) (bool /* completed */, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	maxOffset := frame.Offset + frame.DataLen()
	if err := s.flowController.UpdateHighestReceived(maxOffset, frame.FinBit); err != nil {
		return false, err
	}
	if frame.FinBit {
		s.finalOffset = maxOffset
	}
	if s.canceledRead {
		return frame.FinBit, nil
	}
	if err := s.frameQueue.Push(frame.Data, frame.Offset); err != nil {
		return false, err
	}
	s.signalRead()
	return false, nil
}

func (s *receiveStream) handleResetStreamFrame(frame *wire.ResetStreamFrame) error {
	completed, err := s.handleResetStreamFrameImpl(frame)
	if completed {
		s.streamCompleted()
	}
	return err
}

func (s *receiveStream) handleResetStreamFrameImpl(frame *wire.ResetStreamFrame) (bool /*completed */, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.closedForShutdown {
		return false, nil
	}
	if err := s.flowController.UpdateHighestReceived(frame.ByteOffset, true); err != nil {
		return false, err
	}
	s.finalOffset = frame.ByteOffset

	// ignore duplicate RESET_STREAM frames for this stream (after checking their final offset)
	if s.resetRemotely {
		return false, nil
	}
	s.resetRemotely = true
	s.resetRemotelyErr = streamCanceledError{
		errorCode: frame.ErrorCode,
		error:     fmt.Errorf("Stream %d was reset with error code %d", s.streamID, frame.ErrorCode),
	}
	s.signalRead()
	return true, nil
}

func (s *receiveStream) CloseRemote(offset protocol.ByteCount) {
	s.handleStreamFrame(&wire.StreamFrame{FinBit: true, Offset: offset})
}

func (s *receiveStream) SetReadDeadline(t time.Time) error {
	s.mutex.Lock()
	s.deadline = t
	s.mutex.Unlock()
	s.signalRead()
	return nil
}

// CloseForShutdown closes a stream abruptly.
// It makes Read unblock (and return the error) immediately.
// The peer will NOT be informed about this: the stream is closed without sending a FIN or RESET.
func (s *receiveStream) closeForShutdown(err error) {
	s.mutex.Lock()
	s.closedForShutdown = true
	s.closeForShutdownErr = err
	s.mutex.Unlock()
	s.signalRead()
}

func (s *receiveStream) getWindowUpdate() protocol.ByteCount {
	return s.flowController.GetWindowUpdate()
}

func (s *receiveStream) streamCompleted() {
	s.mutex.Lock()
	finRead := s.finRead
	s.mutex.Unlock()

	if !finRead {
		s.flowController.Abandon()
	}
	s.sender.onStreamCompleted(s.streamID)
}

// signalRead performs a non-blocking send on the readChan
func (s *receiveStream) signalRead() {
	select {
	case s.readChan <- struct{}{}:
	default:
	}
}
<|endoftext|>


<commit_before>package nats

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func startReconnectServer(t *testing.T) *server {
	return startServer(t, 22222, "")
}

func TestReconnectDisallowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = false
	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}

	ts.stopServer()
	if e := wait(ch); e != nil {
		t.Fatal("Did not trigger ClosedCB correctly")
	}
	nc.Close()
}

func TestReconnectAllowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = true
	opts.MaxReconnect = 2
	opts.ReconnectWait = 1 * time.Second

	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ts.stopServer()

	// We want wait to timeout here, and the connection
	// should not trigger the Close CB.
	if e := wait(ch); e == nil {
		t.Fatal("Triggered ClosedCB incorrectly")
	}
	if !nc.isReconnecting() {
		t.Fatal("Expected to be in a reconnecting state")
	}

	// clear the CloseCB since ch will block
	nc.Opts.ClosedCB = nil
	nc.Close()
}

var reconnectOpts = Options{
	Url:            "nats://localhost:22222",
	AllowReconnect: true,
	MaxReconnect:   10,
	ReconnectWait:  100 * time.Millisecond,
	Timeout:        DefaultTimeout,
}

func TestBasicReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestExtendedReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts
	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	rch := make(chan bool)
	opts.ReconnectedCB = func(_ *Conn) {
		rch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}
	testString := "bar"
	received := int32(0)

	ec.Subscribe("foo", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	sub, _ := ec.Subscribe("foobar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	ec.Publish("foo", testString)
	ec.Flush()

	ts.stopServer()
	// server is stopped here..

	// wait for disconnect
	if e := waitTime(dch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a disconnect callback message")
	}

	// Sub while disconnected
	ec.Subscribe("bar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	// Unsub while disconnected
	sub.Unsubscribe()

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	if err = ec.Publish("bar", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	// server is restarted here..
	// wait for reconnect
	if e := waitTime(rch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a reconnect callback message")
	}

	if err = ec.Publish("foobar", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	ch := make(chan bool)
	ec.Subscribe("done", func(b bool) {
		ch <- true
	})
	ec.Publish("done", true)

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	if received != 4 {
		t.Fatalf("Received != %d, equals %d\n", 4, received)
	}
}

func TestParseStateReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	// Simulate partialState, this needs to be cleared
	nc.mu.Lock()
	nc.ps.state = OP_PON
	nc.mu.Unlock()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestQueueSubsOnReconnect(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts

	// Allow us to block on reconnect complete.
	reconnectsDone := make(chan bool)
	opts.ReconnectedCB = func(nc *Conn) {
		reconnectsDone <- true
	}

	// Helper to wait on a reconnect.
	waitOnReconnect := func() {
		select {
		case <-reconnectsDone:
			break
		case <-time.After(2 * time.Second):
			t.Fatalf("Expected a reconnect, timedout!\n")
		}
	}

	// Create connection
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "json")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	// To hold results.
	results := make(map[int]int)
	var mu sync.Mutex

	// Make sure we got what we needed, 1 msg only and all seqnos accounted for..
	checkResults := func(numSent int) {
		mu.Lock()
		defer mu.Unlock()

		for i := 0; i < numSent; i++ {
			if results[i] != 1 {
				t.Fatalf("Received incorrect number of messages, [%d] for seq: %d\n", results[i], i)
			}
		}

		// Auto reset results map
		results = make(map[int]int)
	}

	subj := "foo.bar"
	qgroup := "workers"

	cb := func(seqno int) {
		mu.Lock()
		defer mu.Unlock()
		results[seqno] = results[seqno] + 1
	}

	// Create Queue Subscribers
	ec.QueueSubscribe(subj, qgroup, cb)
	ec.QueueSubscribe(subj, qgroup, cb)

	ec.Flush()

	// Helper function to send messages and check results.
	sendAndCheckMsgs := func(numToSend int) {
		for i := 0; i < numToSend; i++ {
			ec.Publish(subj, i)
		}
		// Wait for processing.
		ec.Flush()
		time.Sleep(50 * time.Millisecond)

		// Check Results
		checkResults(numToSend)
	}

	// Base Test
	sendAndCheckMsgs(10)

	// Stop and restart server
	ts.stopServer()
	ts = startReconnectServer(t)
	defer ts.stopServer()

	waitOnReconnect()

	// Reconnect Base Test
	sendAndCheckMsgs(10)
}

func TestIsClosed(t *testing.T) {
	ts := startReconnectServer(t)
	nc := newConnection(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts.stopServer()
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts = startReconnectServer(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	nc.Close()
	if nc.IsClosed() == false {
		t.Fatalf("IsClosed returned false after Close() was called.")
	}
	ts.stopServer()
}
<commit_msg>Add tests around IsReconnecting and Status helpers<commit_after>package nats

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func startReconnectServer(t *testing.T) *server {
	return startServer(t, 22222, "")
}

func TestReconnectDisallowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = false
	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}

	ts.stopServer()
	if e := wait(ch); e != nil {
		t.Fatal("Did not trigger ClosedCB correctly")
	}
	nc.Close()
}

func TestReconnectAllowedFlags(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = true
	opts.MaxReconnect = 2
	opts.ReconnectWait = 1 * time.Second

	opts.ClosedCB = func(_ *Conn) {
		ch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ts.stopServer()

	// We want wait to timeout here, and the connection
	// should not trigger the Close CB.
	if e := wait(ch); e == nil {
		t.Fatal("Triggered ClosedCB incorrectly")
	}
	if !nc.isReconnecting() {
		t.Fatal("Expected to be in a reconnecting state")
	}

	// clear the CloseCB since ch will block
	nc.Opts.ClosedCB = nil
	nc.Close()
}

var reconnectOpts = Options{
	Url:            "nats://localhost:22222",
	AllowReconnect: true,
	MaxReconnect:   10,
	ReconnectWait:  100 * time.Millisecond,
	Timeout:        DefaultTimeout,
}

func TestBasicReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestExtendedReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts
	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	rch := make(chan bool)
	opts.ReconnectedCB = func(_ *Conn) {
		rch <- true
	}
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}
	testString := "bar"
	received := int32(0)

	ec.Subscribe("foo", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	sub, _ := ec.Subscribe("foobar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	ec.Publish("foo", testString)
	ec.Flush()

	ts.stopServer()
	// server is stopped here..

	// wait for disconnect
	if e := waitTime(dch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a disconnect callback message")
	}

	// Sub while disconnected
	ec.Subscribe("bar", func(s string) {
		atomic.AddInt32(&received, 1)
	})

	// Unsub while disconnected
	sub.Unsubscribe()

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	if err = ec.Publish("bar", testString); err != nil {
		t.Fatalf("Received an error after disconnect: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	// server is restarted here..
	// wait for reconnect
	if e := waitTime(rch, 2*time.Second); e != nil {
		t.Fatal("Did not receive a reconnect callback message")
	}

	if err = ec.Publish("foobar", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	if err = ec.Publish("foo", testString); err != nil {
		t.Fatalf("Received an error after server restarted: %v\n", err)
	}

	ch := make(chan bool)
	ec.Subscribe("done", func(b bool) {
		ch <- true
	})
	ec.Publish("done", true)

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	if received != 4 {
		t.Fatalf("Received != %d, equals %d\n", 4, received)
	}
}

func TestParseStateReconnectFunctionality(t *testing.T) {
	ts := startReconnectServer(t)
	ch := make(chan bool)

	opts := reconnectOpts
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "default")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	testString := "bar"
	ec.Subscribe("foo", func(s string) {
		if s != testString {
			t.Fatal("String doesn't match")
		}
		ch <- true
	})
	ec.Flush()

	// Simulate partialState, this needs to be cleared
	nc.mu.Lock()
	nc.ps.state = OP_PON
	nc.mu.Unlock()

	ts.stopServer()
	// server is stopped here...

	dch := make(chan bool)
	opts.DisconnectedCB = func(_ *Conn) {
		dch <- true
	}
	wait(dch)

	if err := ec.Publish("foo", testString); err != nil {
		t.Fatalf("Failed to publish message: %v\n", err)
	}

	ts = startReconnectServer(t)
	defer ts.stopServer()

	if err := ec.FlushTimeout(5 * time.Second); err != nil {
		t.Fatalf("Error on Flush: %v", err)
	}

	if e := wait(ch); e != nil {
		t.Fatal("Did not receive our message")
	}

	expectedReconnectCount := uint64(1)
	if ec.Conn.Reconnects != expectedReconnectCount {
		t.Fatalf("Reconnect count incorrect: %d vs %d\n",
			ec.Conn.Reconnects, expectedReconnectCount)
	}

	nc.Close()
}

func TestQueueSubsOnReconnect(t *testing.T) {
	ts := startReconnectServer(t)

	opts := reconnectOpts

	// Allow us to block on reconnect complete.
	reconnectsDone := make(chan bool)
	opts.ReconnectedCB = func(nc *Conn) {
		reconnectsDone <- true
	}

	// Helper to wait on a reconnect.
	waitOnReconnect := func() {
		select {
		case <-reconnectsDone:
			break
		case <-time.After(2 * time.Second):
			t.Fatalf("Expected a reconnect, timedout!\n")
		}
	}

	// Create connection
	nc, _ := opts.Connect()
	ec, err := NewEncodedConn(nc, "json")
	if err != nil {
		t.Fatalf("Failed to create an encoded connection: %v\n", err)
	}

	// To hold results.
	results := make(map[int]int)
	var mu sync.Mutex

	// Make sure we got what we needed, 1 msg only and all seqnos accounted for..
	checkResults := func(numSent int) {
		mu.Lock()
		defer mu.Unlock()

		for i := 0; i < numSent; i++ {
			if results[i] != 1 {
				t.Fatalf("Received incorrect number of messages, [%d] for seq: %d\n", results[i], i)
			}
		}

		// Auto reset results map
		results = make(map[int]int)
	}

	subj := "foo.bar"
	qgroup := "workers"

	cb := func(seqno int) {
		mu.Lock()
		defer mu.Unlock()
		results[seqno] = results[seqno] + 1
	}

	// Create Queue Subscribers
	ec.QueueSubscribe(subj, qgroup, cb)
	ec.QueueSubscribe(subj, qgroup, cb)

	ec.Flush()

	// Helper function to send messages and check results.
	sendAndCheckMsgs := func(numToSend int) {
		for i := 0; i < numToSend; i++ {
			ec.Publish(subj, i)
		}
		// Wait for processing.
		ec.Flush()
		time.Sleep(50 * time.Millisecond)

		// Check Results
		checkResults(numToSend)
	}

	// Base Test
	sendAndCheckMsgs(10)

	// Stop and restart server
	ts.stopServer()
	ts = startReconnectServer(t)
	defer ts.stopServer()

	waitOnReconnect()

	// Reconnect Base Test
	sendAndCheckMsgs(10)
}

func TestIsClosed(t *testing.T) {
	ts := startReconnectServer(t)
	nc := newConnection(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts.stopServer()
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	ts = startReconnectServer(t)
	if nc.IsClosed() == true {
		t.Fatalf("IsClosed returned true when the connection is still open.")
	}
	nc.Close()
	if nc.IsClosed() == false {
		t.Fatalf("IsClosed returned false after Close() was called.")
	}
	ts.stopServer()
}

func TestIsReconnectingAndStatus(t *testing.T) {
	ts := startReconnectServer(t)
	// This will kill the last 'ts' server that is created
	defer func() { ts.stopServer() }()
	disconnectedch := make(chan bool)
	reconnectch := make(chan bool)
	opts := DefaultOptions
	opts.Url = "nats://localhost:22222"
	opts.AllowReconnect = true
	opts.MaxReconnect = 10000
	opts.ReconnectWait = 100 * time.Millisecond

	opts.DisconnectedCB = func(_ *Conn) {
		disconnectedch <- true
	}
	opts.ReconnectedCB = func(_ *Conn) {
		reconnectch <- true
	}

	// Connect, verify initial reconnecting state check, then stop the server
	nc, err := opts.Connect()
	if err != nil {
		t.Fatalf("Should have connected ok: %v", err)
	}
	if nc.IsReconnecting() == true {
		t.Fatalf("IsReconnecting returned true when the connection is still open.")
	}
	if status := nc.Status(); status != CONNECTED {
		t.Fatalf("Status returned %d when connected instead of CONNECTED", status)
	}
	ts.stopServer()

	// Wait until we get the disconnected callback
	if e := wait(disconnectedch); e != nil {
		t.Fatalf("Disconnect callback wasn't triggered: %v", e)
	}
	if nc.IsReconnecting() == false {
		t.Fatalf("IsReconnecting returned false when the client is reconnecting.")
	}
	if status := nc.Status(); status != RECONNECTING {
		t.Fatalf("Status returned %d when reconnecting instead of CONNECTED", status)
	}

	ts = startReconnectServer(t)

	// Wait until we get the reconnect callback
	if e := wait(reconnectch); e != nil {
		t.Fatalf("Reconnect callback wasn't triggered: %v", e)
	}
	if nc.IsReconnecting() == true {
		t.Fatalf("IsReconnecting returned true after the connection was reconnected.")
	}
	if status := nc.Status(); status != CONNECTED {
		t.Fatalf("Status returned %d when reconnected instead of CONNECTED", status)
	}

	// Close the connection, reconnecting should still be false
	nc.Close()
	if nc.IsReconnecting() == true {
		t.Fatalf("IsReconnecting returned true after Close() was called.")
	}
	if status := nc.Status(); status != CLOSED {
		t.Fatalf("Status returned %d after Close() was called instead of CLOSED", status)
	}
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"strconv"
	"testing"
	"time"

	"github.com/rafrombrc/gomock/gomock"
	. "github.com/smartystreets/goconvey/convey"
)

func TestBroadcastRouter(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	uaid := "2130ac71-6f04-47cf-b7dc-2570ba1d2afe"
	chid := "90662645-a7b5-4dfe-8105-a290553507e4"
	version := int64(10)
	sentAt := time.Now()

	app := new(Application)
	appConfig := app.ConfigStruct()
	app.Init(nil, appConfig)

	mckLogger := NewMockLogger(mockCtrl)
	app.SetLogger(mckLogger)

	mckStat := NewMockStatistician(mockCtrl)
	app.SetMetrics(mckStat)

	mckStore := NewMockStore(mockCtrl)
	app.SetStore(mckStore)

	mckLocator := NewMockLocator(mockCtrl)

	router := NewBroadcastRouter()
	defaultConfig := router.ConfigStruct()
	conf := defaultConfig.(*BroadcastRouterConfig)
	conf.Listener.Addr = ":8000"
	router.Init(app, conf)
	app.SetRouter(router)

	router.SetLocator(mckLocator)

	srv := new(Serv)
	defSrvConfig := srv.ConfigStruct()
	srvConfig := defSrvConfig.(*ServerConfig)

	srvConfig.Client.Addr = ":10990"
	srvConfig.Endpoint.Addr = ":10899"
	srv.Init(app, defSrvConfig)
	app.SetServer(srv)

	cancelSignal := make(chan bool)

	errChan := make(chan error, 10)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true)
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	go router.Start(errChan)

	Convey("Should fail to route a non-existent uaid", t, func() {

		mckLocator.EXPECT().Contacts(gomock.Any()).Return([]string{}, nil)
		mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).Times(2)
		mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(),
			gomock.Any(), gomock.Any()).Times(2)
		mckStat.EXPECT().Increment("router.broadcast.miss").Times(1)
		mckStat.EXPECT().Timer(gomock.Any(), gomock.Any()).Times(2)
		err := router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")
		So(err, ShouldBeNil)
	})

	Convey("Should succeed self-routing to a valid uaid", t, func() {
		mockWorker := NewMockWorker(mockCtrl)
		client := &Client{mockWorker, &PushWS{}, uaid}
		app.AddClient(uaid, client)

		thisNode := router.URL()
		thisNodeList := []string{thisNode}

		mckLocator.EXPECT().Contacts(gomock.Any()).Return(thisNodeList, nil)
		mckStat.EXPECT().Increment("updates.routed.incoming")
		mckStore.EXPECT().IDsToKey(gomock.Any(), gomock.Any()).Return("", true)
		mckStore.EXPECT().Update(gomock.Any(), gomock.Any()).Return(nil)
		mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
		mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
			gomock.Any()).AnyTimes()
		mockWorker.EXPECT().Flush(gomock.Any(), gomock.Any(), chid, version,
			"").Return(nil)
		mckStat.EXPECT().Gauge("update.client.connections", gomock.Any()).AnyTimes()
		mckStat.EXPECT().Increment("updates.routed.received")
		mckStat.EXPECT().Increment("router.broadcast.hit")
		mckStat.EXPECT().Timer("updates.routed.hits", gomock.Any())
		mckStat.EXPECT().Timer("router.handled", gomock.Any())

		err := router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")
		So(err, ShouldBeNil)
	})

	mckLocator.EXPECT().Close()
	router.Close()
}

func BenchmarkRouter(b *testing.B) {
	for i := 0; i < b.N; i++ {
		runBench(b)
	}
}

func runBench(b *testing.B) {
	mockCtrl := gomock.NewController(b)
	defer mockCtrl.Finish()

	uaid := "2130ac71-6f04-47cf-b7dc-2570ba1d2afe"
	chid := "90662645-a7b5-4dfe-8105-a290553507e4"
	version := int64(10)
	sentAt := time.Now()

	app := new(Application)
	appConfig := app.ConfigStruct()
	app.Init(nil, appConfig)

	mckLogger := NewMockLogger(mockCtrl)
	app.SetLogger(mckLogger)

	mckStat := NewMockStatistician(mockCtrl)
	app.SetMetrics(mckStat)

	mckStore := NewMockStore(mockCtrl)
	app.SetStore(mckStore)

	mckLocator := NewMockLocator(mockCtrl)

	router := NewBroadcastRouter()
	defaultConfig := router.ConfigStruct()
	conf := defaultConfig.(*BroadcastRouterConfig)
	conf.Listener.Addr = ":" + strconv.Itoa(8300+b.N)
	router.Init(app, conf)
	app.SetRouter(router)

	router.SetLocator(mckLocator)

	srv := new(Serv)
	defSrvConfig := srv.ConfigStruct()
	srvConfig := defSrvConfig.(*ServerConfig)

	srvConfig.Client.Addr = ":" + strconv.Itoa(13100+b.N)
	srvConfig.Endpoint.Addr = ":" + strconv.Itoa(13900+b.N)
	mckStat.EXPECT().Gauge("update.client.connections", gomock.Any()).AnyTimes()
	srv.Init(app, defSrvConfig)
	app.SetServer(srv)

	cancelSignal := make(chan bool)

	mockWorker := NewMockWorker(mockCtrl)
	client := &Client{mockWorker, &PushWS{}, uaid}
	app.AddClient(uaid, client)

	thisNode := router.URL()
	thisNodeList := []string{thisNode}

	errChan := make(chan error, 10)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true)
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	go router.Start(errChan)
	<-time.After(time.Duration(1) * time.Second)

	mckLocator.EXPECT().Contacts(gomock.Any()).Return(thisNodeList, nil)
	mckStat.EXPECT().Increment(gomock.Any()).AnyTimes()
	mckStat.EXPECT().Gauge(gomock.Any(), gomock.Any()).AnyTimes()
	mckStat.EXPECT().Timer(gomock.Any(), gomock.Any()).AnyTimes()
	mckStore.EXPECT().IDsToKey(gomock.Any(), gomock.Any()).Return("", true)
	mckStore.EXPECT().Update(gomock.Any(), gomock.Any()).Return(nil)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
		gomock.Any()).AnyTimes()
	mockWorker.EXPECT().Flush(gomock.Any(), gomock.Any(), chid, version,
		"").Return(nil)

	router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")

	mckLocator.EXPECT().Close()
	router.Close()
}
<commit_msg>Add one more test to return an error from a Locator interface.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"errors"
	"strconv"
	"testing"
	"time"

	"github.com/rafrombrc/gomock/gomock"
	. "github.com/smartystreets/goconvey/convey"
)

func TestBroadcastRouter(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	uaid := "2130ac71-6f04-47cf-b7dc-2570ba1d2afe"
	chid := "90662645-a7b5-4dfe-8105-a290553507e4"
	version := int64(10)
	sentAt := time.Now()

	app := new(Application)
	appConfig := app.ConfigStruct()
	app.Init(nil, appConfig)

	mckLogger := NewMockLogger(mockCtrl)
	app.SetLogger(mckLogger)

	mckStat := NewMockStatistician(mockCtrl)
	app.SetMetrics(mckStat)

	mckStore := NewMockStore(mockCtrl)
	app.SetStore(mckStore)

	mckLocator := NewMockLocator(mockCtrl)

	router := NewBroadcastRouter()
	defaultConfig := router.ConfigStruct()
	conf := defaultConfig.(*BroadcastRouterConfig)
	conf.Listener.Addr = ":8000"
	router.Init(app, conf)
	app.SetRouter(router)

	router.SetLocator(mckLocator)

	srv := new(Serv)
	defSrvConfig := srv.ConfigStruct()
	srvConfig := defSrvConfig.(*ServerConfig)

	srvConfig.Client.Addr = ":10990"
	srvConfig.Endpoint.Addr = ":10899"
	srv.Init(app, defSrvConfig)
	app.SetServer(srv)

	cancelSignal := make(chan bool)

	errChan := make(chan error, 10)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true)
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	go router.Start(errChan)

	Convey("Should fail to route a non-existent uaid", t, func() {

		mckLocator.EXPECT().Contacts(gomock.Any()).Return([]string{}, nil)
		mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).Times(2)
		mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(),
			gomock.Any(), gomock.Any()).Times(2)
		mckStat.EXPECT().Increment("router.broadcast.miss").Times(1)
		mckStat.EXPECT().Timer(gomock.Any(), gomock.Any()).Times(2)
		err := router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")
		So(err, ShouldBeNil)
	})

	Convey("Should fail to route if contacts errors", t, func() {
		myErr := errors.New("Oops")
		mckLocator.EXPECT().Contacts(gomock.Any()).Return([]string{}, myErr)
		mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).Times(2)
		mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(),
			gomock.Any(), gomock.Any()).Times(2)
		mckStat.EXPECT().Increment("router.broadcast.error").Times(1)
		err := router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")
		So(err, ShouldEqual, myErr)
	})

	Convey("Should succeed self-routing to a valid uaid", t, func() {
		mockWorker := NewMockWorker(mockCtrl)
		client := &Client{mockWorker, &PushWS{}, uaid}
		app.AddClient(uaid, client)

		thisNode := router.URL()
		thisNodeList := []string{thisNode}

		mckLocator.EXPECT().Contacts(gomock.Any()).Return(thisNodeList, nil)
		mckStat.EXPECT().Increment("updates.routed.incoming")
		mckStore.EXPECT().IDsToKey(gomock.Any(), gomock.Any()).Return("", true)
		mckStore.EXPECT().Update(gomock.Any(), gomock.Any()).Return(nil)
		mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
		mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
			gomock.Any()).AnyTimes()
		mockWorker.EXPECT().Flush(gomock.Any(), gomock.Any(), chid, version,
			"").Return(nil)
		mckStat.EXPECT().Gauge("update.client.connections", gomock.Any()).AnyTimes()
		mckStat.EXPECT().Increment("updates.routed.received")
		mckStat.EXPECT().Increment("router.broadcast.hit")
		mckStat.EXPECT().Timer("updates.routed.hits", gomock.Any())
		mckStat.EXPECT().Timer("router.handled", gomock.Any())

		err := router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")
		So(err, ShouldBeNil)
	})

	mckLocator.EXPECT().Close()
	router.Close()
}

func BenchmarkRouter(b *testing.B) {
	for i := 0; i < b.N; i++ {
		runBench(b)
	}
}

func runBench(b *testing.B) {
	mockCtrl := gomock.NewController(b)
	defer mockCtrl.Finish()

	uaid := "2130ac71-6f04-47cf-b7dc-2570ba1d2afe"
	chid := "90662645-a7b5-4dfe-8105-a290553507e4"
	version := int64(10)
	sentAt := time.Now()

	app := new(Application)
	appConfig := app.ConfigStruct()
	app.Init(nil, appConfig)

	mckLogger := NewMockLogger(mockCtrl)
	app.SetLogger(mckLogger)

	mckStat := NewMockStatistician(mockCtrl)
	app.SetMetrics(mckStat)

	mckStore := NewMockStore(mockCtrl)
	app.SetStore(mckStore)

	mckLocator := NewMockLocator(mockCtrl)

	router := NewBroadcastRouter()
	defaultConfig := router.ConfigStruct()
	conf := defaultConfig.(*BroadcastRouterConfig)
	conf.Listener.Addr = ":" + strconv.Itoa(8300+b.N)
	router.Init(app, conf)
	app.SetRouter(router)

	router.SetLocator(mckLocator)

	srv := new(Serv)
	defSrvConfig := srv.ConfigStruct()
	srvConfig := defSrvConfig.(*ServerConfig)

	srvConfig.Client.Addr = ":" + strconv.Itoa(13100+b.N)
	srvConfig.Endpoint.Addr = ":" + strconv.Itoa(13900+b.N)
	mckStat.EXPECT().Gauge("update.client.connections", gomock.Any()).AnyTimes()
	srv.Init(app, defSrvConfig)
	app.SetServer(srv)

	cancelSignal := make(chan bool)

	mockWorker := NewMockWorker(mockCtrl)
	client := &Client{mockWorker, &PushWS{}, uaid}
	app.AddClient(uaid, client)

	thisNode := router.URL()
	thisNodeList := []string{thisNode}

	errChan := make(chan error, 10)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true)
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	go router.Start(errChan)
	<-time.After(time.Duration(1) * time.Second)

	mckLocator.EXPECT().Contacts(gomock.Any()).Return(thisNodeList, nil)
	mckStat.EXPECT().Increment(gomock.Any()).AnyTimes()
	mckStat.EXPECT().Gauge(gomock.Any(), gomock.Any()).AnyTimes()
	mckStat.EXPECT().Timer(gomock.Any(), gomock.Any()).AnyTimes()
	mckStore.EXPECT().IDsToKey(gomock.Any(), gomock.Any()).Return("", true)
	mckStore.EXPECT().Update(gomock.Any(), gomock.Any()).Return(nil)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
		gomock.Any()).AnyTimes()
	mockWorker.EXPECT().Flush(gomock.Any(), gomock.Any(), chid, version,
		"").Return(nil)

	router.Route(cancelSignal, uaid, chid, version, sentAt, "", "")

	mckLocator.EXPECT().Close()
	router.Close()
}
<|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 hostgw

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/coreos/flannel/Godeps/_workspace/src/github.com/golang/glog"
	"github.com/coreos/flannel/Godeps/_workspace/src/github.com/vishvananda/netlink"
	"github.com/coreos/flannel/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/coreos/flannel/backend"
	"github.com/coreos/flannel/pkg/ip"
	"github.com/coreos/flannel/subnet"
)

const (
	routeCheckRetries = 10
)

type HostgwBackend struct {
	sm       subnet.Manager
	network  string
	lease    *subnet.Lease
	extIface *net.Interface
	extIP    net.IP
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
	rl       []netlink.Route
}

func New(sm subnet.Manager, network string) backend.Backend {
	ctx, cancel := context.WithCancel(context.Background())

	b := &HostgwBackend{
		sm:      sm,
		network: network,
		ctx:     ctx,
		cancel:  cancel,
	}
	return b
}

func (rb *HostgwBackend) Init(extIface *net.Interface, extIP net.IP) (*backend.SubnetDef, error) {
	rb.extIface = extIface
	rb.extIP = extIP

	attrs := subnet.LeaseAttrs{
		PublicIP:    ip.FromIP(extIP),
		BackendType: "host-gw",
	}

	l, err := rb.sm.AcquireLease(rb.ctx, rb.network, &attrs)
	switch err {
	case nil:
		rb.lease = l

	case context.Canceled, context.DeadlineExceeded:
		return nil, err

	default:
		return nil, fmt.Errorf("failed to acquire lease: %v", err)
	}

	/* NB: docker will create the local route to `sn` */

	return &backend.SubnetDef{
		Net: l.Subnet,
		MTU: extIface.MTU,
	}, nil
}

func (rb *HostgwBackend) Run() {
	rb.wg.Add(1)
	go func() {
		subnet.LeaseRenewer(rb.ctx, rb.sm, rb.network, rb.lease)
		rb.wg.Done()
	}()

	log.Info("Watching for new subnet leases")
	evts := make(chan []subnet.Event)
	rb.wg.Add(1)
	go func() {
		subnet.WatchLeases(rb.ctx, rb.sm, rb.network, evts)
		rb.wg.Done()
	}()

	rb.rl = make([]netlink.Route, 0, 10)
	rb.wg.Add(1)
	go func() {
		rb.routeCheck(rb.ctx)
		rb.wg.Done()
	}()

	defer rb.wg.Wait()

	for {
		select {
		case evtBatch := <-evts:
			rb.handleSubnetEvents(evtBatch)

		case <-rb.ctx.Done():
			return
		}
	}
}

func (rb *HostgwBackend) Stop() {
	rb.cancel()
}

func (rb *HostgwBackend) Name() string {
	return "host-gw"
}

func (rb *HostgwBackend) handleSubnetEvents(batch []subnet.Event) {
	for _, evt := range batch {
		switch evt.Type {
		case subnet.SubnetAdded:
			log.Infof("Subnet added: %v via %v", evt.Lease.Subnet, evt.Lease.Attrs.PublicIP)

			if evt.Lease.Attrs.BackendType != "host-gw" {
				log.Warningf("Ignoring non-host-gw subnet: type=%v", evt.Lease.Attrs.BackendType)
				continue
			}

			route := netlink.Route{
				Dst:       evt.Lease.Subnet.ToIPNet(),
				Gw:        evt.Lease.Attrs.PublicIP.ToIP(),
				LinkIndex: rb.extIface.Index,
			}
			if err := netlink.RouteAdd(&route); err != nil {
				log.Errorf("Error adding route to %v via %v: %v", evt.Lease.Subnet, evt.Lease.Attrs.PublicIP, err)
				continue
			}
			rb.addToRouteList(route)

		case subnet.SubnetRemoved:
			log.Info("Subnet removed: ", evt.Lease.Subnet)

			if evt.Lease.Attrs.BackendType != "host-gw" {
				log.Warningf("Ignoring non-host-gw subnet: type=%v", evt.Lease.Attrs.BackendType)
				continue
			}

			route := netlink.Route{
				Dst:       evt.Lease.Subnet.ToIPNet(),
				Gw:        evt.Lease.Attrs.PublicIP.ToIP(),
				LinkIndex: rb.extIface.Index,
			}
			if err := netlink.RouteDel(&route); err != nil {
				log.Errorf("Error deleting route to %v: %v", evt.Lease.Subnet, err)
				continue
			}
			rb.removeFromRouteList(route)

		default:
			log.Error("Internal error: unknown event type: ", int(evt.Type))
		}
	}
}

func (rb *HostgwBackend) addToRouteList(route netlink.Route) {
	rb.rl = append(rb.rl, route)
}

func (rb *HostgwBackend) removeFromRouteList(route netlink.Route) {
	for index, r := range rb.rl {
		if routeEqual(r, route) {
			rb.rl = append(rb.rl[:index], rb.rl[index+1:]...)
			return
		}
	}
}

func (rb *HostgwBackend) routeCheck(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(routeCheckRetries * time.Second):
			rb.checkSubnetExistInRoutes()
		}
	}
}

func (rb *HostgwBackend) checkSubnetExistInRoutes() {
	routeList, err := netlink.RouteList(nil, netlink.FAMILY_V4)
	if err == nil {
		for _, route := range rb.rl {
			exist := false
			for _, r := range routeList {
				if r.Dst == nil {
					continue
				}
				if routeEqual(r, route) {
					exist = true
					break
				}
			}
			if !exist {
				if err := netlink.RouteAdd(&route); err != nil {
					if nerr, ok := err.(net.Error); !ok {
						log.Errorf("Error recovering route to %v: %v, %v", route.Dst, route.Gw, nerr)
					}
					continue
				} else {
					log.Infof("Route recovered %v : %v", route.Dst, route.Gw)
				}
			}
		}
	}
}

func routeEqual(x, y netlink.Route) bool {
	if x.Dst.IP.Equal(y.Dst.IP) && x.Gw.Equal(y.Gw) && bytes.Equal(x.Dst.Mask, y.Dst.Mask) {
		return true
	}
	return false
}
<commit_msg>Ignore route to local ext IP in hostgw mode<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 hostgw

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/coreos/flannel/Godeps/_workspace/src/github.com/golang/glog"
	"github.com/coreos/flannel/Godeps/_workspace/src/github.com/vishvananda/netlink"
	"github.com/coreos/flannel/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/coreos/flannel/backend"
	"github.com/coreos/flannel/pkg/ip"
	"github.com/coreos/flannel/subnet"
)

const (
	routeCheckRetries = 10
)

type HostgwBackend struct {
	sm       subnet.Manager
	network  string
	lease    *subnet.Lease
	extIface *net.Interface
	extIP    net.IP
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
	rl       []netlink.Route
}

func New(sm subnet.Manager, network string) backend.Backend {
	ctx, cancel := context.WithCancel(context.Background())

	b := &HostgwBackend{
		sm:      sm,
		network: network,
		ctx:     ctx,
		cancel:  cancel,
	}
	return b
}

func (rb *HostgwBackend) Init(extIface *net.Interface, extIP net.IP) (*backend.SubnetDef, error) {
	rb.extIface = extIface
	rb.extIP = extIP

	attrs := subnet.LeaseAttrs{
		PublicIP:    ip.FromIP(extIP),
		BackendType: "host-gw",
	}

	l, err := rb.sm.AcquireLease(rb.ctx, rb.network, &attrs)
	switch err {
	case nil:
		rb.lease = l

	case context.Canceled, context.DeadlineExceeded:
		return nil, err

	default:
		return nil, fmt.Errorf("failed to acquire lease: %v", err)
	}

	/* NB: docker will create the local route to `sn` */

	return &backend.SubnetDef{
		Net: l.Subnet,
		MTU: extIface.MTU,
	}, nil
}

func (rb *HostgwBackend) Run() {
	rb.wg.Add(1)
	go func() {
		subnet.LeaseRenewer(rb.ctx, rb.sm, rb.network, rb.lease)
		rb.wg.Done()
	}()

	log.Info("Watching for new subnet leases")
	evts := make(chan []subnet.Event)
	rb.wg.Add(1)
	go func() {
		subnet.WatchLeases(rb.ctx, rb.sm, rb.network, evts)
		rb.wg.Done()
	}()

	rb.rl = make([]netlink.Route, 0, 10)
	rb.wg.Add(1)
	go func() {
		rb.routeCheck(rb.ctx)
		rb.wg.Done()
	}()

	defer rb.wg.Wait()

	for {
		select {
		case evtBatch := <-evts:
			rb.handleSubnetEvents(evtBatch)

		case <-rb.ctx.Done():
			return
		}
	}
}

func (rb *HostgwBackend) Stop() {
	rb.cancel()
}

func (rb *HostgwBackend) Name() string {
	return "host-gw"
}

func (rb *HostgwBackend) handleSubnetEvents(batch []subnet.Event) {
	for _, evt := range batch {
		switch evt.Type {
		case subnet.SubnetAdded:
			log.Infof("Subnet added: %v via %v", evt.Lease.Subnet, evt.Lease.Attrs.PublicIP)

			if evt.Lease.Attrs.BackendType != "host-gw" {
				log.Warningf("Ignoring non-host-gw subnet: type=%v", evt.Lease.Attrs.BackendType)
				continue
			}

			route := netlink.Route{
				Dst:       evt.Lease.Subnet.ToIPNet(),
				Gw:        evt.Lease.Attrs.PublicIP.ToIP(),
				LinkIndex: rb.extIface.Index,
			}
			if rb.extIP.Equal(route.Gw) {
				continue
			}
			if err := netlink.RouteAdd(&route); err != nil {
				log.Errorf("Error adding route to %v via %v: %v", evt.Lease.Subnet, evt.Lease.Attrs.PublicIP, err)
				continue
			}
			rb.addToRouteList(route)

		case subnet.SubnetRemoved:
			log.Info("Subnet removed: ", evt.Lease.Subnet)

			if evt.Lease.Attrs.BackendType != "host-gw" {
				log.Warningf("Ignoring non-host-gw subnet: type=%v", evt.Lease.Attrs.BackendType)
				continue
			}

			route := netlink.Route{
				Dst:       evt.Lease.Subnet.ToIPNet(),
				Gw:        evt.Lease.Attrs.PublicIP.ToIP(),
				LinkIndex: rb.extIface.Index,
			}
			if err := netlink.RouteDel(&route); err != nil {
				log.Errorf("Error deleting route to %v: %v", evt.Lease.Subnet, err)
				continue
			}
			rb.removeFromRouteList(route)

		default:
			log.Error("Internal error: unknown event type: ", int(evt.Type))
		}
	}
}

func (rb *HostgwBackend) addToRouteList(route netlink.Route) {
	rb.rl = append(rb.rl, route)
}

func (rb *HostgwBackend) removeFromRouteList(route netlink.Route) {
	for index, r := range rb.rl {
		if routeEqual(r, route) {
			rb.rl = append(rb.rl[:index], rb.rl[index+1:]...)
			return
		}
	}
}

func (rb *HostgwBackend) routeCheck(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(routeCheckRetries * time.Second):
			rb.checkSubnetExistInRoutes()
		}
	}
}

func (rb *HostgwBackend) checkSubnetExistInRoutes() {
	routeList, err := netlink.RouteList(nil, netlink.FAMILY_V4)
	if err == nil {
		for _, route := range rb.rl {
			exist := false
			for _, r := range routeList {
				if r.Dst == nil {
					continue
				}
				if routeEqual(r, route) {
					exist = true
					break
				}
			}
			if !exist {
				if err := netlink.RouteAdd(&route); err != nil {
					if nerr, ok := err.(net.Error); !ok {
						log.Errorf("Error recovering route to %v: %v, %v", route.Dst, route.Gw, nerr)
					}
					continue
				} else {
					log.Infof("Route recovered %v : %v", route.Dst, route.Gw)
				}
			}
		}
	}
}

func routeEqual(x, y netlink.Route) bool {
	if x.Dst.IP.Equal(y.Dst.IP) && x.Gw.Equal(y.Gw) && bytes.Equal(x.Dst.Mask, y.Dst.Mask) {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package vulcand

import (
	"errors"
	"github.com/3onyc/hipdate/backends"
	"github.com/3onyc/hipdate/shared"
	vulcan "github.com/mailgun/vulcand/api"
	"github.com/mailgun/vulcand/plugin/registry"
	"log"
	"net/http"
	"strings"
)

var (
	MissingApiUrlError = errors.New("vulcand api endpoint not specified")
)

type VulcandBackend struct {
	v *vulcan.Client
}

func NewVulcandBackend(opts shared.OptionMap) (backends.Backend, error) {
	eu, ok := opts["url"]
	if !ok {
		return nil, MissingApiUrlError
	}

	v, err := createClient(eu)
	if err != nil {
		return nil, err
	}

	return &VulcandBackend{
		v: v,
	}, nil
}

func (vb *VulcandBackend) AddUpstream(
	h shared.Host,
	u shared.Upstream,
) error {
	hName := string(h)
	uUrl := string(u)
	uId := hName + "_up"
	eId := hName + "_ep_" + u.Hash()
	lId := hName + "_loc"

	if _, err := vb.v.AddHost(string(h)); isError(err) {
		return err
	}

	if _, err := vb.v.AddUpstream(uId); isError(err) {
		return err
	}

	if _, err := vb.v.AddEndpoint(uId, eId, uUrl); isError(err) {
		return err
	}

	if _, err := vb.v.AddLocation(hName, lId, "/.*", uId); isError(err) {
		return err
	}

	return nil
}
func (vb *VulcandBackend) RemoveUpstream(
	h shared.Host,
	u shared.Upstream,
) error {
	uId := string(h) + "_up"
	eId := string(h) + "_ep_" + u.Hash()

	if _, err := vb.v.DeleteEndpoint(uId, eId); isError(err) {
		return err
	}

	return nil
}

func (vb *VulcandBackend) Initialise() error {
	hosts, err := vb.v.GetHosts()
	if err != nil {
		return err
	}

	for _, h := range hosts {
		if _, err := vb.v.DeleteHost(h.Name); isError(err) {
			return err
		}
	}

	upstreams, err := vb.v.GetUpstreams()
	if err != nil {
		return err
	}

	for _, u := range upstreams {
		if _, err := vb.v.DeleteUpstream(u.Id); isError(err) {
			return err
		}
	}

	return nil
}

func (vb *VulcandBackend) ListHosts() (*shared.HostList, error) {
	hl := shared.HostList{}

	hs, err := vb.v.GetHosts()
	if err != nil {
		return nil, err
	}

	for _, vh := range hs {
		if len(vh.Locations) == 0 {
			continue
		}

		l := vh.Locations[0]
		h := shared.Host(l.Hostname)

		hl[h] = []shared.Upstream{}

		if l.Upstream != nil && len(l.Upstream.Endpoints) > 0 {
			for _, ep := range l.Upstream.Endpoints {
				hl[h] = append(hl[h], shared.Upstream(ep.Url))
			}
		}
	}

	return &hl, nil
}

func isError(err error) bool {
	if err == nil {
		return false
	}

	log.Printf("DEBUG [backend:vulcand] isError: %s\n", err)

	if strings.Contains(err.Error(), "already exists") {
		return false
	}

	if strings.Contains(err.Error(), "already exists") {
		return false
	}

	return true
}

func init() {
	backends.BackendMap["vulcand"] = NewVulcandBackend
}

func createClient(eu string) (*vulcan.Client, error) {
	v := vulcan.NewClient(eu, registry.GetRegistry())

	// Check if etcd is reachable
	if _, err := http.Get(eu + "/v1/"); err != nil {
		return nil, err
	}

	return v, nil
}
<commit_msg>Remove isError debug in vulcand backend<commit_after>package vulcand

import (
	"errors"
	"github.com/3onyc/hipdate/backends"
	"github.com/3onyc/hipdate/shared"
	vulcan "github.com/mailgun/vulcand/api"
	"github.com/mailgun/vulcand/plugin/registry"
	"net/http"
	"strings"
)

var (
	MissingApiUrlError = errors.New("vulcand api endpoint not specified")
)

type VulcandBackend struct {
	v *vulcan.Client
}

func NewVulcandBackend(opts shared.OptionMap) (backends.Backend, error) {
	eu, ok := opts["url"]
	if !ok {
		return nil, MissingApiUrlError
	}

	v, err := createClient(eu)
	if err != nil {
		return nil, err
	}

	return &VulcandBackend{
		v: v,
	}, nil
}

func (vb *VulcandBackend) AddUpstream(
	h shared.Host,
	u shared.Upstream,
) error {
	hName := string(h)
	uUrl := string(u)
	uId := hName + "_up"
	eId := hName + "_ep_" + u.Hash()
	lId := hName + "_loc"

	if _, err := vb.v.AddHost(string(h)); isError(err) {
		return err
	}

	if _, err := vb.v.AddUpstream(uId); isError(err) {
		return err
	}

	if _, err := vb.v.AddEndpoint(uId, eId, uUrl); isError(err) {
		return err
	}

	if _, err := vb.v.AddLocation(hName, lId, "/.*", uId); isError(err) {
		return err
	}

	return nil
}
func (vb *VulcandBackend) RemoveUpstream(
	h shared.Host,
	u shared.Upstream,
) error {
	uId := string(h) + "_up"
	eId := string(h) + "_ep_" + u.Hash()

	if _, err := vb.v.DeleteEndpoint(uId, eId); isError(err) {
		return err
	}

	return nil
}

func (vb *VulcandBackend) Initialise() error {
	hosts, err := vb.v.GetHosts()
	if err != nil {
		return err
	}

	for _, h := range hosts {
		if _, err := vb.v.DeleteHost(h.Name); isError(err) {
			return err
		}
	}

	upstreams, err := vb.v.GetUpstreams()
	if err != nil {
		return err
	}

	for _, u := range upstreams {
		if _, err := vb.v.DeleteUpstream(u.Id); isError(err) {
			return err
		}
	}

	return nil
}

func (vb *VulcandBackend) ListHosts() (*shared.HostList, error) {
	hl := shared.HostList{}

	hs, err := vb.v.GetHosts()
	if err != nil {
		return nil, err
	}

	for _, vh := range hs {
		if len(vh.Locations) == 0 {
			continue
		}

		l := vh.Locations[0]
		h := shared.Host(l.Hostname)

		hl[h] = []shared.Upstream{}

		if l.Upstream != nil && len(l.Upstream.Endpoints) > 0 {
			for _, ep := range l.Upstream.Endpoints {
				hl[h] = append(hl[h], shared.Upstream(ep.Url))
			}
		}
	}

	return &hl, nil
}

func isError(err error) bool {
	if err == nil {
		return false
	}

	if strings.Contains(err.Error(), "already exists") {
		return false
	}

	if strings.Contains(err.Error(), "already exists") {
		return false
	}

	return true
}

func init() {
	backends.BackendMap["vulcand"] = NewVulcandBackend
}

func createClient(eu string) (*vulcan.Client, error) {
	v := vulcan.NewClient(eu, registry.GetRegistry())

	// Check if etcd is reachable
	if _, err := http.Get(eu + "/v1/"); err != nil {
		return nil, err
	}

	return v, nil
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
	"encoding/hex"
	"fmt"
	capn "github.com/glycerine/go-capnproto"
	"sort"
)

const (
	RMIdEmpty = RMId(0)
	KeyLen    = 20
	ClientLen = KeyLen - 8
)

type RMId uint32

func (rmId RMId) String() string {
	if rmId == RMIdEmpty {
		return "RM:Empty"
	} else {
		return fmt.Sprintf("RM:%x", uint32(rmId))
	}
}

type KeyType [KeyLen]byte

type VarUUId KeyType

func MakeVarUUId(data []byte) *VarUUId {
	uuid := VarUUId([KeyLen]byte{})
	copy(uuid[:], data)
	return &uuid
}

func MakeVarUUIdFromStr(str string) *VarUUId {
	data, err := hex.DecodeString(str)
	if err != nil {
		return nil
	}
	return MakeVarUUId(data)
}

func (vUUId VarUUId) String() string {
	h := ([]byte)(hex.EncodeToString(vUUId[:]))
	return fmt.Sprintf("VarUUId:%s", h)
}

type TxnId KeyType

func MakeTxnId(data []byte) *TxnId {
	id := TxnId([KeyLen]byte{})
	copy(id[:], data)
	return &id
}

func (txnId TxnId) String() string {
	t := txnId[:]
	client, conn, boot, rmId := hex.EncodeToString(t[0:8]), hex.EncodeToString(t[8:12]), hex.EncodeToString(t[12:16]), hex.EncodeToString(t[16:20])
	return fmt.Sprintf("TxnId:%s-%s-%s-%v", client, conn, boot, rmId)
}

func (txnId TxnId) ClientId() [ClientLen]byte {
	client := [ClientLen]byte{}
	copy(client[:], txnId[8:])
	return client
}

func (a *TxnId) Equal(b *TxnId) bool {
	return a == b || (a != nil && b != nil && bytes.Equal(a[:], b[:]))
}

func (a *TxnId) LessThan(b *TxnId) bool {
	switch {
	case b == nil:
		return false
	case a == nil:
		return true
	default:
		return bytes.Compare(a[:], b[:]) < 0
	}
}

func (a *VarUUId) Equal(b *VarUUId) bool {
	return a == b || (a != nil && b != nil && bytes.Equal(a[:], b[:]))
}

func (a *VarUUId) LessThan(b *VarUUId) bool {
	switch {
	case b == nil:
		return false
	case a == nil:
		return true
	default:
		return bytes.Compare(a[:], b[:]) < 0
	}
}

type VarUUIds []*VarUUId

func (vUUIds VarUUIds) Sort()              { sort.Sort(vUUIds) }
func (vUUIds VarUUIds) Len() int           { return len(vUUIds) }
func (vUUIds VarUUIds) Less(i, j int) bool { return vUUIds[i].LessThan(vUUIds[j]) }
func (vUUIds VarUUIds) Swap(i, j int)      { vUUIds[i], vUUIds[j] = vUUIds[j], vUUIds[i] }

type Positions capn.UInt8List

func (a *Positions) Equal(b *Positions) bool {
	if a == nil || b == nil {
		return a == b
	}
	aCap, bCap := (*capn.UInt8List)(a), (*capn.UInt8List)(b)
	if aCap.Len() != bCap.Len() {
		return false
	}
	for idx, l := 0, aCap.Len(); idx < l; idx++ {
		if aV, bV := aCap.At(idx), bCap.At(idx); aV != bV {
			return false
		}
	}
	return true
}

func (p *Positions) String() string {
	return fmt.Sprint((*capn.UInt8List)(p).ToArray())
}

type RMIds []RMId

func (rmIds RMIds) Sort()              { sort.Sort(rmIds) }
func (rmIds RMIds) Len() int           { return len(rmIds) }
func (rmIds RMIds) Less(i, j int) bool { return rmIds[i] < rmIds[j] }
func (rmIds RMIds) Swap(i, j int)      { rmIds[i], rmIds[j] = rmIds[j], rmIds[i] }

func (a RMIds) Equal(b RMIds) bool {
	if a == nil || b == nil {
		return a == nil && b == nil
	}
	if len(a) != len(b) {
		return false
	}
	for idx, rA := range a {
		if b[idx] != rA {
			return false
		}
	}
	return true
}

func (rmIds RMIds) NonEmptyLen() int {
	count := 0
	for _, rmId := range rmIds {
		if rmId != RMIdEmpty {
			count++
		}
	}
	return count
}
<commit_msg>WIP Ref T6.<commit_after>package common

import (
	"bytes"
	"encoding/hex"
	"fmt"
	capn "github.com/glycerine/go-capnproto"
	"sort"
)

const (
	RMIdEmpty = RMId(0)
	KeyLen    = 20
	ClientLen = KeyLen - 8
)

type RMId uint32

func (rmId RMId) String() string {
	if rmId == RMIdEmpty {
		return "RM:Empty"
	} else {
		return fmt.Sprintf("RM:%x", uint32(rmId))
	}
}

type KeyType [KeyLen]byte

type VarUUId KeyType

func MakeVarUUId(data []byte) *VarUUId {
	uuid := VarUUId([KeyLen]byte{})
	copy(uuid[:], data)
	return &uuid
}

func MakeVarUUIdFromStr(str string) *VarUUId {
	data, err := hex.DecodeString(str)
	if err != nil {
		return nil
	}
	return MakeVarUUId(data)
}

func (vUUId VarUUId) String() string {
	h := ([]byte)(hex.EncodeToString(vUUId[:]))
	return fmt.Sprintf("VarUUId:%s", h)
}

type TxnId KeyType

func MakeTxnId(data []byte) *TxnId {
	id := TxnId([KeyLen]byte{})
	copy(id[:], data)
	return &id
}

func (txnId TxnId) String() string {
	t := txnId[:]
	client, conn, boot, rmId := hex.EncodeToString(t[0:8]), hex.EncodeToString(t[8:12]), hex.EncodeToString(t[12:16]), hex.EncodeToString(t[16:20])
	return fmt.Sprintf("TxnId:%s-%s-%s-%v", client, conn, boot, rmId)
}

func (txnId TxnId) ClientId() [ClientLen]byte {
	client := [ClientLen]byte{}
	copy(client[:], txnId[8:])
	return client
}

func (a *TxnId) Equal(b *TxnId) bool {
	return a == b || (a != nil && b != nil && bytes.Equal(a[:], b[:]))
}

func (a *TxnId) LessThan(b *TxnId) bool {
	switch {
	case b == nil:
		return false
	case a == nil:
		return true
	default:
		return bytes.Compare(a[:], b[:]) < 0
	}
}

func (a *VarUUId) Equal(b *VarUUId) bool {
	return a == b || (a != nil && b != nil && bytes.Equal(a[:], b[:]))
}

func (a *VarUUId) LessThan(b *VarUUId) bool {
	switch {
	case b == nil:
		return false
	case a == nil:
		return true
	default:
		return bytes.Compare(a[:], b[:]) < 0
	}
}

type VarUUIds []*VarUUId

func (vUUIds VarUUIds) Sort()              { sort.Sort(vUUIds) }
func (vUUIds VarUUIds) Len() int           { return len(vUUIds) }
func (vUUIds VarUUIds) Less(i, j int) bool { return vUUIds[i].LessThan(vUUIds[j]) }
func (vUUIds VarUUIds) Swap(i, j int)      { vUUIds[i], vUUIds[j] = vUUIds[j], vUUIds[i] }

type Positions capn.UInt8List

func (a *Positions) Equal(b *Positions) bool {
	if a == nil || b == nil {
		return a == b
	}
	aCap, bCap := (*capn.UInt8List)(a), (*capn.UInt8List)(b)
	if aCap.Len() != bCap.Len() {
		return false
	}
	for idx, l := 0, aCap.Len(); idx < l; idx++ {
		if aV, bV := aCap.At(idx), bCap.At(idx); aV != bV {
			return false
		}
	}
	return true
}

func (p *Positions) String() string {
	return fmt.Sprint((*capn.UInt8List)(p).ToArray())
}

type RMIds []RMId

func (rmIds RMIds) Sort()              { sort.Sort(rmIds) }
func (rmIds RMIds) Len() int           { return len(rmIds) }
func (rmIds RMIds) Less(i, j int) bool { return rmIds[i] < rmIds[j] }
func (rmIds RMIds) Swap(i, j int)      { rmIds[i], rmIds[j] = rmIds[j], rmIds[i] }

func (a RMIds) Equal(b RMIds) bool {
	if a == nil || b == nil {
		return a == nil && b == nil
	}
	if len(a) != len(b) {
		return false
	}
	for idx, rA := range a {
		if b[idx] != rA {
			return false
		}
	}
	return true
}

func (rmIds RMIds) NonEmptyLen() int {
	count := 0
	for _, rmId := range rmIds {
		if rmId != RMIdEmpty {
			count++
		}
	}
	return count
}

func (rmIds RMIds) NonEmpty() RMIds {
	nel := rmIds.NonEmptyLen()
	if nel == len(rmIds) {
		return rmIds
	}
	nonEmpty := make([]RMId, 0, nel)
	for _, rmId := range rmIds {
		if rmId != RMIdEmpty {
			nonEmpty = append(nonEmpty, rmId)
		}
	}
	return nonEmpty
}
<|endoftext|>


<commit_before>
<commit_msg>base implementation 02<commit_after><|endoftext|>