<commit_before>// Copyright 2018 The Prometheus 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 wal

import (
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/pkg/errors"

	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
	"github.com/prometheus/prometheus/tsdb/fileutil"
	"github.com/prometheus/prometheus/tsdb/record"
	"github.com/prometheus/prometheus/tsdb/tombstones"
)

// CheckpointStats returns stats about a created checkpoint.
type CheckpointStats struct {
	DroppedSeries     int
	DroppedSamples    int
	DroppedTombstones int
	DroppedExemplars  int
	TotalSeries       int // Processed series including dropped ones.
	TotalSamples      int // Processed samples including dropped ones.
	TotalTombstones   int // Processed tombstones including dropped ones.
	TotalExemplars    int // Processed exemplars including dropped ones.
}

// LastCheckpoint returns the directory name and index of the most recent checkpoint.
// If dir does not contain any checkpoints, ErrNotFound is returned.
func LastCheckpoint(dir string) (string, int, error) {
	checkpoints, err := listCheckpoints(dir)
	if err != nil {
		return "", 0, err
	}

	if len(checkpoints) == 0 {
		return "", 0, record.ErrNotFound
	}

	checkpoint := checkpoints[len(checkpoints)-1]
	return filepath.Join(dir, checkpoint.name), checkpoint.index, nil
}

// DeleteCheckpoints deletes all checkpoints in a directory below a given index.
func DeleteCheckpoints(dir string, maxIndex int) error {
	checkpoints, err := listCheckpoints(dir)
	if err != nil {
		return err
	}

	errs := tsdb_errors.NewMulti()
	for _, checkpoint := range checkpoints {
		if checkpoint.index >= maxIndex {
			break
		}
		errs.Add(os.RemoveAll(filepath.Join(dir, checkpoint.name)))
	}
	return errs.Err()
}

const checkpointPrefix = "checkpoint."

// Checkpoint creates a compacted checkpoint of segments in range [first, last] in the given WAL.
// It includes the most recent checkpoint if it exists.
// All series not satisfying keep and samples below mint are dropped.
//
// The checkpoint is stored in a directory named checkpoint.N in the same
// segmented format as the original WAL itself.
// This makes it easy to read it through the WAL package and concatenate
// it with the original WAL.
func Checkpoint(logger log.Logger, w *WAL, from, to int, keep func(id uint64) bool, mint int64) (*CheckpointStats, error) {
	stats := &CheckpointStats{}
	var sgmReader io.ReadCloser

	level.Info(logger).Log("msg", "Creating checkpoint", "from_segment", from, "to_segment", to, "mint", mint)

	{

		var sgmRange []SegmentRange
		dir, idx, err := LastCheckpoint(w.Dir())
		if err != nil && err != record.ErrNotFound {
			return nil, errors.Wrap(err, "find last checkpoint")
		}
		last := idx + 1
		if err == nil {
			if from > last {
				return nil, fmt.Errorf("unexpected gap to last checkpoint. expected:%v, requested:%v", last, from)
			}
			// Ignore WAL files below the checkpoint. They shouldn't exist to begin with.
			from = last

			sgmRange = append(sgmRange, SegmentRange{Dir: dir, Last: math.MaxInt32})
		}

		sgmRange = append(sgmRange, SegmentRange{Dir: w.Dir(), First: from, Last: to})
		sgmReader, err = NewSegmentsRangeReader(sgmRange...)
		if err != nil {
			return nil, errors.Wrap(err, "create segment reader")
		}
		defer sgmReader.Close()
	}

	cpdir := checkpointDir(w.Dir(), to)
	cpdirtmp := cpdir + ".tmp"

	if err := os.RemoveAll(cpdirtmp); err != nil {
		return nil, errors.Wrap(err, "remove previous temporary checkpoint dir")
	}

	if err := os.MkdirAll(cpdirtmp, 0o777); err != nil {
		return nil, errors.Wrap(err, "create checkpoint dir")
	}
	cp, err := New(nil, nil, cpdirtmp, w.CompressionEnabled())
	if err != nil {
		return nil, errors.Wrap(err, "open checkpoint")
	}

	// Ensures that an early return caused by an error doesn't leave any tmp files.
	defer func() {
		cp.Close()
		os.RemoveAll(cpdirtmp)
	}()

	r := NewReader(sgmReader)

	var (
		series    []record.RefSeries
		samples   []record.RefSample
		tstones   []tombstones.Stone
		exemplars []record.RefExemplar
		dec       record.Decoder
		enc       record.Encoder
		buf       []byte
		recs      [][]byte
	)
	for r.Next() {
		series, samples, tstones, exemplars = series[:0], samples[:0], tstones[:0], exemplars[:0]

		// We don't reset the buffer since we batch up multiple records
		// before writing them to the checkpoint.
		// Remember where the record for this iteration starts.
		start := len(buf)
		rec := r.Record()

		switch dec.Type(rec) {
		case record.Series:
			series, err = dec.Series(rec, series)
			if err != nil {
				return nil, errors.Wrap(err, "decode series")
			}
			// Drop irrelevant series in place.
			repl := series[:0]
			for _, s := range series {
				if keep(s.Ref) {
					repl = append(repl, s)
				}
			}
			if len(repl) > 0 {
				buf = enc.Series(repl, buf)
			}
			stats.TotalSeries += len(series)
			stats.DroppedSeries += len(series) - len(repl)

		case record.Samples:
			samples, err = dec.Samples(rec, samples)
			if err != nil {
				return nil, errors.Wrap(err, "decode samples")
			}
			// Drop irrelevant samples in place.
			repl := samples[:0]
			for _, s := range samples {
				if s.T >= mint {
					repl = append(repl, s)
				}
			}
			if len(repl) > 0 {
				buf = enc.Samples(repl, buf)
			}
			stats.TotalSamples += len(samples)
			stats.DroppedSamples += len(samples) - len(repl)

		case record.Tombstones:
			tstones, err = dec.Tombstones(rec, tstones)
			if err != nil {
				return nil, errors.Wrap(err, "decode deletes")
			}
			// Drop irrelevant tombstones in place.
			repl := tstones[:0]
			for _, s := range tstones {
				for _, iv := range s.Intervals {
					if iv.Maxt >= mint {
						repl = append(repl, s)
						break
					}
				}
			}
			if len(repl) > 0 {
				buf = enc.Tombstones(repl, buf)
			}
			stats.TotalTombstones += len(tstones)
			stats.DroppedTombstones += len(tstones) - len(repl)

		case record.Exemplars:
			exemplars, err = dec.Exemplars(rec, exemplars)
			if err != nil {
				return nil, errors.Wrap(err, "decode exemplars")
			}
			// Drop irrelevant exemplars in place.
			repl := exemplars[:0]
			for _, e := range exemplars {
				if e.T >= mint {
					repl = append(repl, e)
				}
			}
			if len(repl) > 0 {
				buf = enc.Exemplars(repl, buf)
			}
			stats.TotalExemplars += len(exemplars)
			stats.DroppedExemplars += len(exemplars) - len(repl)
		default:
			// Unknown record type, probably from a future Prometheus version.
			continue
		}
		if len(buf[start:]) == 0 {
			continue // All contents discarded.
		}
		recs = append(recs, buf[start:])

		// Flush records in 1 MB increments.
		if len(buf) > 1*1024*1024 {
			if err := cp.Log(recs...); err != nil {
				return nil, errors.Wrap(err, "flush records")
			}
			buf, recs = buf[:0], recs[:0]
		}
	}
	// If we hit any corruption during checkpointing, repairing is not an option.
	// The head won't know which series records are lost.
	if r.Err() != nil {
		return nil, errors.Wrap(r.Err(), "read segments")
	}

	// Flush remaining records.
	if err := cp.Log(recs...); err != nil {
		return nil, errors.Wrap(err, "flush records")
	}
	if err := cp.Close(); err != nil {
		return nil, errors.Wrap(err, "close checkpoint")
	}

	// Sync temporary directory before rename.
	df, err := fileutil.OpenDir(cpdirtmp)
	if err != nil {
		return nil, errors.Wrap(err, "open temporary checkpoint directory")
	}
	if err := df.Sync(); err != nil {
		df.Close()
		return nil, errors.Wrap(err, "sync temporary checkpoint directory")
	}
	if err = df.Close(); err != nil {
		return nil, errors.Wrap(err, "close temporary checkpoint directory")
	}

	if err := fileutil.Replace(cpdirtmp, cpdir); err != nil {
		return nil, errors.Wrap(err, "rename checkpoint directory")
	}

	return stats, nil
}

func checkpointDir(dir string, i int) string {
	return filepath.Join(dir, fmt.Sprintf(checkpointPrefix+"%08d", i))
}

type checkpointRef struct {
	name  string
	index int
}

func listCheckpoints(dir string) (refs []checkpointRef, err error) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return nil, err
	}

	for i := 0; i < len(files); i++ {
		fi := files[i]
		if !strings.HasPrefix(fi.Name(), checkpointPrefix) {
			continue
		}
		if !fi.IsDir() {
			return nil, errors.Errorf("checkpoint %s is not a directory", fi.Name())
		}
		idx, err := strconv.Atoi(fi.Name()[len(checkpointPrefix):])
		if err != nil {
			continue
		}

		refs = append(refs, checkpointRef{name: fi.Name(), index: idx})
	}

	sort.Slice(refs, func(i, j int) bool {
		return refs[i].index < refs[j].index
	})

	return refs, nil
}
<commit_msg>refine comments of Checkpoint function (#9655)<commit_after>// Copyright 2018 The Prometheus 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 wal

import (
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/pkg/errors"

	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
	"github.com/prometheus/prometheus/tsdb/fileutil"
	"github.com/prometheus/prometheus/tsdb/record"
	"github.com/prometheus/prometheus/tsdb/tombstones"
)

// CheckpointStats returns stats about a created checkpoint.
type CheckpointStats struct {
	DroppedSeries     int
	DroppedSamples    int
	DroppedTombstones int
	DroppedExemplars  int
	TotalSeries       int // Processed series including dropped ones.
	TotalSamples      int // Processed samples including dropped ones.
	TotalTombstones   int // Processed tombstones including dropped ones.
	TotalExemplars    int // Processed exemplars including dropped ones.
}

// LastCheckpoint returns the directory name and index of the most recent checkpoint.
// If dir does not contain any checkpoints, ErrNotFound is returned.
func LastCheckpoint(dir string) (string, int, error) {
	checkpoints, err := listCheckpoints(dir)
	if err != nil {
		return "", 0, err
	}

	if len(checkpoints) == 0 {
		return "", 0, record.ErrNotFound
	}

	checkpoint := checkpoints[len(checkpoints)-1]
	return filepath.Join(dir, checkpoint.name), checkpoint.index, nil
}

// DeleteCheckpoints deletes all checkpoints in a directory below a given index.
func DeleteCheckpoints(dir string, maxIndex int) error {
	checkpoints, err := listCheckpoints(dir)
	if err != nil {
		return err
	}

	errs := tsdb_errors.NewMulti()
	for _, checkpoint := range checkpoints {
		if checkpoint.index >= maxIndex {
			break
		}
		errs.Add(os.RemoveAll(filepath.Join(dir, checkpoint.name)))
	}
	return errs.Err()
}

const checkpointPrefix = "checkpoint."

// Checkpoint creates a compacted checkpoint of segments in range [from, to] in the given WAL.
// It includes the most recent checkpoint if it exists.
// All series not satisfying keep and samples/tombstones/exemplars below mint are dropped.
//
// The checkpoint is stored in a directory named checkpoint.N in the same
// segmented format as the original WAL itself.
// This makes it easy to read it through the WAL package and concatenate
// it with the original WAL.
func Checkpoint(logger log.Logger, w *WAL, from, to int, keep func(id uint64) bool, mint int64) (*CheckpointStats, error) {
	stats := &CheckpointStats{}
	var sgmReader io.ReadCloser

	level.Info(logger).Log("msg", "Creating checkpoint", "from_segment", from, "to_segment", to, "mint", mint)

	{
		var sgmRange []SegmentRange
		dir, idx, err := LastCheckpoint(w.Dir())
		if err != nil && err != record.ErrNotFound {
			return nil, errors.Wrap(err, "find last checkpoint")
		}
		last := idx + 1
		if err == nil {
			if from > last {
				return nil, fmt.Errorf("unexpected gap to last checkpoint. expected:%v, requested:%v", last, from)
			}
			// Ignore WAL files below the checkpoint. They shouldn't exist to begin with.
			from = last

			sgmRange = append(sgmRange, SegmentRange{Dir: dir, Last: math.MaxInt32})
		}

		sgmRange = append(sgmRange, SegmentRange{Dir: w.Dir(), First: from, Last: to})
		sgmReader, err = NewSegmentsRangeReader(sgmRange...)
		if err != nil {
			return nil, errors.Wrap(err, "create segment reader")
		}
		defer sgmReader.Close()
	}

	cpdir := checkpointDir(w.Dir(), to)
	cpdirtmp := cpdir + ".tmp"

	if err := os.RemoveAll(cpdirtmp); err != nil {
		return nil, errors.Wrap(err, "remove previous temporary checkpoint dir")
	}

	if err := os.MkdirAll(cpdirtmp, 0o777); err != nil {
		return nil, errors.Wrap(err, "create checkpoint dir")
	}
	cp, err := New(nil, nil, cpdirtmp, w.CompressionEnabled())
	if err != nil {
		return nil, errors.Wrap(err, "open checkpoint")
	}

	// Ensures that an early return caused by an error doesn't leave any tmp files.
	defer func() {
		cp.Close()
		os.RemoveAll(cpdirtmp)
	}()

	r := NewReader(sgmReader)

	var (
		series    []record.RefSeries
		samples   []record.RefSample
		tstones   []tombstones.Stone
		exemplars []record.RefExemplar
		dec       record.Decoder
		enc       record.Encoder
		buf       []byte
		recs      [][]byte
	)
	for r.Next() {
		series, samples, tstones, exemplars = series[:0], samples[:0], tstones[:0], exemplars[:0]

		// We don't reset the buffer since we batch up multiple records
		// before writing them to the checkpoint.
		// Remember where the record for this iteration starts.
		start := len(buf)
		rec := r.Record()

		switch dec.Type(rec) {
		case record.Series:
			series, err = dec.Series(rec, series)
			if err != nil {
				return nil, errors.Wrap(err, "decode series")
			}
			// Drop irrelevant series in place.
			repl := series[:0]
			for _, s := range series {
				if keep(s.Ref) {
					repl = append(repl, s)
				}
			}
			if len(repl) > 0 {
				buf = enc.Series(repl, buf)
			}
			stats.TotalSeries += len(series)
			stats.DroppedSeries += len(series) - len(repl)

		case record.Samples:
			samples, err = dec.Samples(rec, samples)
			if err != nil {
				return nil, errors.Wrap(err, "decode samples")
			}
			// Drop irrelevant samples in place.
			repl := samples[:0]
			for _, s := range samples {
				if s.T >= mint {
					repl = append(repl, s)
				}
			}
			if len(repl) > 0 {
				buf = enc.Samples(repl, buf)
			}
			stats.TotalSamples += len(samples)
			stats.DroppedSamples += len(samples) - len(repl)

		case record.Tombstones:
			tstones, err = dec.Tombstones(rec, tstones)
			if err != nil {
				return nil, errors.Wrap(err, "decode deletes")
			}
			// Drop irrelevant tombstones in place.
			repl := tstones[:0]
			for _, s := range tstones {
				for _, iv := range s.Intervals {
					if iv.Maxt >= mint {
						repl = append(repl, s)
						break
					}
				}
			}
			if len(repl) > 0 {
				buf = enc.Tombstones(repl, buf)
			}
			stats.TotalTombstones += len(tstones)
			stats.DroppedTombstones += len(tstones) - len(repl)

		case record.Exemplars:
			exemplars, err = dec.Exemplars(rec, exemplars)
			if err != nil {
				return nil, errors.Wrap(err, "decode exemplars")
			}
			// Drop irrelevant exemplars in place.
			repl := exemplars[:0]
			for _, e := range exemplars {
				if e.T >= mint {
					repl = append(repl, e)
				}
			}
			if len(repl) > 0 {
				buf = enc.Exemplars(repl, buf)
			}
			stats.TotalExemplars += len(exemplars)
			stats.DroppedExemplars += len(exemplars) - len(repl)
		default:
			// Unknown record type, probably from a future Prometheus version.
			continue
		}
		if len(buf[start:]) == 0 {
			continue // All contents discarded.
		}
		recs = append(recs, buf[start:])

		// Flush records in 1 MB increments.
		if len(buf) > 1*1024*1024 {
			if err := cp.Log(recs...); err != nil {
				return nil, errors.Wrap(err, "flush records")
			}
			buf, recs = buf[:0], recs[:0]
		}
	}
	// If we hit any corruption during checkpointing, repairing is not an option.
	// The head won't know which series records are lost.
	if r.Err() != nil {
		return nil, errors.Wrap(r.Err(), "read segments")
	}

	// Flush remaining records.
	if err := cp.Log(recs...); err != nil {
		return nil, errors.Wrap(err, "flush records")
	}
	if err := cp.Close(); err != nil {
		return nil, errors.Wrap(err, "close checkpoint")
	}

	// Sync temporary directory before rename.
	df, err := fileutil.OpenDir(cpdirtmp)
	if err != nil {
		return nil, errors.Wrap(err, "open temporary checkpoint directory")
	}
	if err := df.Sync(); err != nil {
		df.Close()
		return nil, errors.Wrap(err, "sync temporary checkpoint directory")
	}
	if err = df.Close(); err != nil {
		return nil, errors.Wrap(err, "close temporary checkpoint directory")
	}

	if err := fileutil.Replace(cpdirtmp, cpdir); err != nil {
		return nil, errors.Wrap(err, "rename checkpoint directory")
	}

	return stats, nil
}

func checkpointDir(dir string, i int) string {
	return filepath.Join(dir, fmt.Sprintf(checkpointPrefix+"%08d", i))
}

type checkpointRef struct {
	name  string
	index int
}

func listCheckpoints(dir string) (refs []checkpointRef, err error) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return nil, err
	}

	for i := 0; i < len(files); i++ {
		fi := files[i]
		if !strings.HasPrefix(fi.Name(), checkpointPrefix) {
			continue
		}
		if !fi.IsDir() {
			return nil, errors.Errorf("checkpoint %s is not a directory", fi.Name())
		}
		idx, err := strconv.Atoi(fi.Name()[len(checkpointPrefix):])
		if err != nil {
			continue
		}

		refs = append(refs, checkpointRef{name: fi.Name(), index: idx})
	}

	sort.Slice(refs, func(i, j int) bool {
		return refs[i].index < refs[j].index
	})

	return refs, nil
}
<|endoftext|>


<commit_before>package tsdb

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/influxdb/models"
)

func NewStore(path string) *Store {
	opts := NewEngineOptions()
	opts.Config = NewConfig()

	return &Store{
		path:          path,
		EngineOptions: opts,
		Logger:        log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

var (
	ErrShardNotFound = fmt.Errorf("shard not found")
)

const (
	MaintenanceCheckInterval = time.Minute
)

type Store struct {
	mu   sync.RWMutex
	path string

	databaseIndexes map[string]*DatabaseIndex
	shards          map[uint64]*Shard

	EngineOptions EngineOptions
	Logger        *log.Logger

	closing chan struct{}
	wg      sync.WaitGroup
}

// Path returns the store's root path.
func (s *Store) Path() string { return s.path }

// DatabaseIndexN returns the number of databases indicies in the store.
func (s *Store) DatabaseIndexN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.databaseIndexes)
}

// Shard returns a shard by id.
func (s *Store) Shard(id uint64) *Shard {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.shards[id]
}

// ShardN returns the number of shard in the store.
func (s *Store) ShardN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.shards)
}

func (s *Store) CreateShard(database, retentionPolicy string, shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.closing:
		return fmt.Errorf("closing")
	default:
	}

	// shard already exists
	if _, ok := s.shards[shardID]; ok {
		return nil
	}

	// created the db and retention policy dirs if they don't exist
	if err := os.MkdirAll(filepath.Join(s.path, database, retentionPolicy), 0700); err != nil {
		return err
	}

	// create the WAL directory
	walPath := filepath.Join(s.EngineOptions.Config.WALDir, database, retentionPolicy, fmt.Sprintf("%d", shardID))
	if err := os.MkdirAll(walPath, 0700); err != nil {
		return err
	}

	// create the database index if it does not exist
	db, ok := s.databaseIndexes[database]
	if !ok {
		db = NewDatabaseIndex()
		s.databaseIndexes[database] = db
	}

	shardPath := filepath.Join(s.path, database, retentionPolicy, strconv.FormatUint(shardID, 10))
	shard := NewShard(shardID, db, shardPath, walPath, s.EngineOptions)
	if err := shard.Open(); err != nil {
		return err
	}

	s.shards[shardID] = shard

	return nil
}

// DeleteShard removes a shard from disk.
func (s *Store) DeleteShard(shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// ensure shard exists
	sh, ok := s.shards[shardID]
	if !ok {
		return nil
	}

	if err := sh.Close(); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.path); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.walPath); err != nil {
		return err
	}

	delete(s.shards, shardID)

	return nil
}

// DeleteDatabase will close all shards associated with a database and remove the directory and files from disk.
func (s *Store) DeleteDatabase(name string, shardIDs []uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, id := range shardIDs {
		shard := s.shards[id]
		if shard != nil {
			shard.Close()
		}
	}
	if err := os.RemoveAll(filepath.Join(s.path, name)); err != nil {
		return err
	}
	if err := os.RemoveAll(filepath.Join(s.EngineOptions.Config.WALDir, name)); err != nil {
		return err
	}
	delete(s.databaseIndexes, name)
	return nil
}

// ShardIDs returns a slice of all ShardIDs under management.
func (s *Store) ShardIDs() []uint64 {
	ids := make([]uint64, 0, len(s.shards))
	for i, _ := range s.shards {
		ids = append(ids, i)
	}
	return ids
}

func (s *Store) ValidateAggregateFieldsInStatement(shardID uint64, measurementName string, stmt *influxql.SelectStatement) error {
	s.mu.RLock()
	shard := s.shards[shardID]
	s.mu.RUnlock()
	if shard == nil {
		return ErrShardNotFound
	}
	return shard.ValidateAggregateFieldsInStatement(measurementName, stmt)
}

func (s *Store) DatabaseIndex(name string) *DatabaseIndex {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.databaseIndexes[name]
}

// Databases returns all the databases in the indexes
func (s *Store) Databases() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	databases := []string{}
	for db := range s.databaseIndexes {
		databases = append(databases, db)
	}
	return databases
}

func (s *Store) Measurement(database, name string) *Measurement {
	s.mu.RLock()
	db := s.databaseIndexes[database]
	s.mu.RUnlock()
	if db == nil {
		return nil
	}
	return db.Measurement(name)
}

// DiskSize returns the size of all the shard files in bytes.  This size does not include the WAL size.
func (s *Store) DiskSize() (int64, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	var size int64
	for _, shardID := range s.ShardIDs() {
		shard := s.Shard(shardID)
		sz, err := shard.DiskSize()
		if err != nil {
			return 0, err
		}
		size += sz
	}
	return size, nil
}

// deleteSeries loops through the local shards and deletes the series data and metadata for the passed in series keys
func (s *Store) deleteSeries(keys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteSeries(keys); err != nil {
			return err
		}
	}
	return nil
}

// deleteMeasurement loops through the local shards and removes the measurement field encodings from each shard
func (s *Store) deleteMeasurement(name string, seriesKeys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteMeasurement(name, seriesKeys); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) loadIndexes() error {
	dbs, err := ioutil.ReadDir(s.path)
	if err != nil {
		return err
	}
	for _, db := range dbs {
		if !db.IsDir() {
			s.Logger.Printf("Skipping database dir: %s. Not a directory", db.Name())
			continue
		}
		s.databaseIndexes[db.Name()] = NewDatabaseIndex()
	}
	return nil
}

func (s *Store) loadShards() error {
	// loop through the current database indexes
	for db := range s.databaseIndexes {
		rps, err := ioutil.ReadDir(filepath.Join(s.path, db))
		if err != nil {
			return err
		}

		for _, rp := range rps {
			// retention policies should be directories.  Skip anything that is not a dir.
			if !rp.IsDir() {
				s.Logger.Printf("Skipping retention policy dir: %s. Not a directory", rp.Name())
				continue
			}

			shards, err := ioutil.ReadDir(filepath.Join(s.path, db, rp.Name()))
			if err != nil {
				return err
			}
			for _, sh := range shards {
				path := filepath.Join(s.path, db, rp.Name(), sh.Name())
				walPath := filepath.Join(s.EngineOptions.Config.WALDir, db, rp.Name(), sh.Name())

				// Shard file names are numeric shardIDs
				shardID, err := strconv.ParseUint(sh.Name(), 10, 64)
				if err != nil {
					s.Logger.Printf("Skipping shard: %s. Not a valid path", rp.Name())
					continue
				}

				shard := NewShard(shardID, s.databaseIndexes[db], path, walPath, s.EngineOptions)
				err = shard.Open()
				if err != nil {
					return fmt.Errorf("failed to open shard %d: %s", shardID, err)
				}
				s.shards[shardID] = shard
			}
		}
	}
	return nil

}

// periodicMaintenance is the method called in a goroutine on the opening of the store
// to perform periodic maintenance of the shards.
func (s *Store) periodicMaintenance() {
	t := time.NewTicker(MaintenanceCheckInterval)
	for {
		select {
		case <-t.C:
			s.performMaintenance()
		case <-s.closing:
			t.Stop()
			return
		}
	}
}

// performMaintenance will loop through the shars and tell them
// to perform any maintenance tasks. Those tasks should kick off
// their own goroutines if it's anything that could take time.
func (s *Store) performMaintenance() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, sh := range s.shards {
		s.performMaintenanceOnShard(sh)
	}
}

func (s *Store) performMaintenanceOnShard(shard *Shard) {
	defer func() {
		if r := recover(); r != nil {
			s.Logger.Printf("recovered eror in maintenance on shard %d", shard.id)
		}
	}()
	shard.PerformMaintenance()
}

func (s *Store) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.closing = make(chan struct{})

	s.shards = map[uint64]*Shard{}
	s.databaseIndexes = map[string]*DatabaseIndex{}

	s.Logger.Printf("Using data dir: %v", s.Path())

	// Create directory.
	if err := os.MkdirAll(s.path, 0777); err != nil {
		return err
	}

	// TODO: Start AE for Node
	if err := s.loadIndexes(); err != nil {
		return err
	}

	if err := s.loadShards(); err != nil {
		return err
	}

	go s.periodicMaintenance()

	return nil
}

func (s *Store) WriteToShard(shardID uint64, points []models.Point) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sh, ok := s.shards[shardID]
	if !ok {
		return ErrShardNotFound
	}

	return sh.WritePoints(points)
}

func (s *Store) CreateMapper(shardID uint64, stmt influxql.Statement, chunkSize int) (Mapper, error) {
	shard := s.Shard(shardID)

	switch stmt := stmt.(type) {
	case *influxql.SelectStatement:
		if (stmt.IsRawQuery && !stmt.HasDistinct()) || stmt.IsSimpleDerivative() {
			m := NewRawMapper(shard, stmt)
			m.ChunkSize = chunkSize
			return m, nil
		}
		return NewAggregateMapper(shard, stmt), nil

	case *influxql.ShowMeasurementsStatement:
		m := NewShowMeasurementsMapper(shard, stmt)
		m.ChunkSize = chunkSize
		return m, nil
	case *influxql.ShowTagKeysStatement:
		return NewShowTagKeysMapper(shard, stmt, chunkSize), nil
	default:
		return nil, fmt.Errorf("can't create mapper for statement type: %T", stmt)
	}
}

func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closing != nil {
		close(s.closing)
		s.closing = nil
	}
	s.wg.Wait()

	for _, sh := range s.shards {
		if err := sh.Close(); err != nil {
			return err
		}
	}
	if s.closing != nil {
		close(s.closing)
	}
	s.closing = nil
	s.shards = nil
	s.databaseIndexes = nil

	return nil
}

// IsRetryable returns true if this error is temporary and could be retried
func IsRetryable(err error) bool {
	if err == nil {
		return true
	}

	if strings.Contains(err.Error(), "field type conflict") {
		return false
	}
	return true
}
<commit_msg>Fix panic: assignment to entry in nil map<commit_after>package tsdb

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/influxdb/models"
)

func NewStore(path string) *Store {
	opts := NewEngineOptions()
	opts.Config = NewConfig()

	return &Store{
		path:          path,
		EngineOptions: opts,
		Logger:        log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

var (
	ErrShardNotFound = fmt.Errorf("shard not found")
	ErrStoreClosed   = fmt.Errorf("store is closed")
)

const (
	MaintenanceCheckInterval = time.Minute
)

type Store struct {
	mu   sync.RWMutex
	path string

	databaseIndexes map[string]*DatabaseIndex
	shards          map[uint64]*Shard

	EngineOptions EngineOptions
	Logger        *log.Logger

	closing chan struct{}
	wg      sync.WaitGroup
	opened  bool
}

// Path returns the store's root path.
func (s *Store) Path() string { return s.path }

// DatabaseIndexN returns the number of databases indicies in the store.
func (s *Store) DatabaseIndexN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.databaseIndexes)
}

// Shard returns a shard by id.
func (s *Store) Shard(id uint64) *Shard {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.shards[id]
}

// ShardN returns the number of shard in the store.
func (s *Store) ShardN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.shards)
}

func (s *Store) CreateShard(database, retentionPolicy string, shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.closing:
		return ErrStoreClosed
	default:
	}

	// shard already exists
	if _, ok := s.shards[shardID]; ok {
		return nil
	}

	// created the db and retention policy dirs if they don't exist
	if err := os.MkdirAll(filepath.Join(s.path, database, retentionPolicy), 0700); err != nil {
		return err
	}

	// create the WAL directory
	walPath := filepath.Join(s.EngineOptions.Config.WALDir, database, retentionPolicy, fmt.Sprintf("%d", shardID))
	if err := os.MkdirAll(walPath, 0700); err != nil {
		return err
	}

	// create the database index if it does not exist
	db, ok := s.databaseIndexes[database]
	if !ok {
		db = NewDatabaseIndex()
		s.databaseIndexes[database] = db
	}

	shardPath := filepath.Join(s.path, database, retentionPolicy, strconv.FormatUint(shardID, 10))
	shard := NewShard(shardID, db, shardPath, walPath, s.EngineOptions)
	if err := shard.Open(); err != nil {
		return err
	}

	s.shards[shardID] = shard

	return nil
}

// DeleteShard removes a shard from disk.
func (s *Store) DeleteShard(shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// ensure shard exists
	sh, ok := s.shards[shardID]
	if !ok {
		return nil
	}

	if err := sh.Close(); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.path); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.walPath); err != nil {
		return err
	}

	delete(s.shards, shardID)

	return nil
}

// DeleteDatabase will close all shards associated with a database and remove the directory and files from disk.
func (s *Store) DeleteDatabase(name string, shardIDs []uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, id := range shardIDs {
		shard := s.shards[id]
		if shard != nil {
			shard.Close()
		}
	}
	if err := os.RemoveAll(filepath.Join(s.path, name)); err != nil {
		return err
	}
	if err := os.RemoveAll(filepath.Join(s.EngineOptions.Config.WALDir, name)); err != nil {
		return err
	}
	delete(s.databaseIndexes, name)
	return nil
}

// ShardIDs returns a slice of all ShardIDs under management.
func (s *Store) ShardIDs() []uint64 {
	ids := make([]uint64, 0, len(s.shards))
	for i, _ := range s.shards {
		ids = append(ids, i)
	}
	return ids
}

func (s *Store) ValidateAggregateFieldsInStatement(shardID uint64, measurementName string, stmt *influxql.SelectStatement) error {
	s.mu.RLock()
	shard := s.shards[shardID]
	s.mu.RUnlock()
	if shard == nil {
		return ErrShardNotFound
	}
	return shard.ValidateAggregateFieldsInStatement(measurementName, stmt)
}

func (s *Store) DatabaseIndex(name string) *DatabaseIndex {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.databaseIndexes[name]
}

// Databases returns all the databases in the indexes
func (s *Store) Databases() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	databases := []string{}
	for db := range s.databaseIndexes {
		databases = append(databases, db)
	}
	return databases
}

func (s *Store) Measurement(database, name string) *Measurement {
	s.mu.RLock()
	db := s.databaseIndexes[database]
	s.mu.RUnlock()
	if db == nil {
		return nil
	}
	return db.Measurement(name)
}

// DiskSize returns the size of all the shard files in bytes.  This size does not include the WAL size.
func (s *Store) DiskSize() (int64, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	var size int64
	for _, shardID := range s.ShardIDs() {
		shard := s.Shard(shardID)
		sz, err := shard.DiskSize()
		if err != nil {
			return 0, err
		}
		size += sz
	}
	return size, nil
}

// deleteSeries loops through the local shards and deletes the series data and metadata for the passed in series keys
func (s *Store) deleteSeries(keys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteSeries(keys); err != nil {
			return err
		}
	}
	return nil
}

// deleteMeasurement loops through the local shards and removes the measurement field encodings from each shard
func (s *Store) deleteMeasurement(name string, seriesKeys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteMeasurement(name, seriesKeys); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) loadIndexes() error {
	dbs, err := ioutil.ReadDir(s.path)
	if err != nil {
		return err
	}
	for _, db := range dbs {
		if !db.IsDir() {
			s.Logger.Printf("Skipping database dir: %s. Not a directory", db.Name())
			continue
		}
		s.databaseIndexes[db.Name()] = NewDatabaseIndex()
	}
	return nil
}

func (s *Store) loadShards() error {
	// loop through the current database indexes
	for db := range s.databaseIndexes {
		rps, err := ioutil.ReadDir(filepath.Join(s.path, db))
		if err != nil {
			return err
		}

		for _, rp := range rps {
			// retention policies should be directories.  Skip anything that is not a dir.
			if !rp.IsDir() {
				s.Logger.Printf("Skipping retention policy dir: %s. Not a directory", rp.Name())
				continue
			}

			shards, err := ioutil.ReadDir(filepath.Join(s.path, db, rp.Name()))
			if err != nil {
				return err
			}
			for _, sh := range shards {
				path := filepath.Join(s.path, db, rp.Name(), sh.Name())
				walPath := filepath.Join(s.EngineOptions.Config.WALDir, db, rp.Name(), sh.Name())

				// Shard file names are numeric shardIDs
				shardID, err := strconv.ParseUint(sh.Name(), 10, 64)
				if err != nil {
					s.Logger.Printf("Skipping shard: %s. Not a valid path", rp.Name())
					continue
				}

				shard := NewShard(shardID, s.databaseIndexes[db], path, walPath, s.EngineOptions)
				err = shard.Open()
				if err != nil {
					return fmt.Errorf("failed to open shard %d: %s", shardID, err)
				}
				s.shards[shardID] = shard
			}
		}
	}
	return nil

}

// periodicMaintenance is the method called in a goroutine on the opening of the store
// to perform periodic maintenance of the shards.
func (s *Store) periodicMaintenance() {
	t := time.NewTicker(MaintenanceCheckInterval)
	for {
		select {
		case <-t.C:
			s.performMaintenance()
		case <-s.closing:
			t.Stop()
			return
		}
	}
}

// performMaintenance will loop through the shars and tell them
// to perform any maintenance tasks. Those tasks should kick off
// their own goroutines if it's anything that could take time.
func (s *Store) performMaintenance() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, sh := range s.shards {
		s.performMaintenanceOnShard(sh)
	}
}

func (s *Store) performMaintenanceOnShard(shard *Shard) {
	defer func() {
		if r := recover(); r != nil {
			s.Logger.Printf("recovered eror in maintenance on shard %d", shard.id)
		}
	}()
	shard.PerformMaintenance()
}

func (s *Store) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.closing = make(chan struct{})

	s.shards = map[uint64]*Shard{}
	s.databaseIndexes = map[string]*DatabaseIndex{}

	s.Logger.Printf("Using data dir: %v", s.Path())

	// Create directory.
	if err := os.MkdirAll(s.path, 0777); err != nil {
		return err
	}

	// TODO: Start AE for Node
	if err := s.loadIndexes(); err != nil {
		return err
	}

	if err := s.loadShards(); err != nil {
		return err
	}

	go s.periodicMaintenance()
	s.opened = true

	return nil
}

func (s *Store) WriteToShard(shardID uint64, points []models.Point) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	select {
	case <-s.closing:
		return ErrStoreClosed
	default:
	}

	sh, ok := s.shards[shardID]
	if !ok {
		return ErrShardNotFound
	}

	return sh.WritePoints(points)
}

func (s *Store) CreateMapper(shardID uint64, stmt influxql.Statement, chunkSize int) (Mapper, error) {
	shard := s.Shard(shardID)

	switch stmt := stmt.(type) {
	case *influxql.SelectStatement:
		if (stmt.IsRawQuery && !stmt.HasDistinct()) || stmt.IsSimpleDerivative() {
			m := NewRawMapper(shard, stmt)
			m.ChunkSize = chunkSize
			return m, nil
		}
		return NewAggregateMapper(shard, stmt), nil

	case *influxql.ShowMeasurementsStatement:
		m := NewShowMeasurementsMapper(shard, stmt)
		m.ChunkSize = chunkSize
		return m, nil
	case *influxql.ShowTagKeysStatement:
		return NewShowTagKeysMapper(shard, stmt, chunkSize), nil
	default:
		return nil, fmt.Errorf("can't create mapper for statement type: %T", stmt)
	}
}

func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.opened {
		close(s.closing)
	}
	s.wg.Wait()

	for _, sh := range s.shards {
		if err := sh.Close(); err != nil {
			return err
		}
	}
	s.opened = false
	s.shards = nil
	s.databaseIndexes = nil

	return nil
}

// IsRetryable returns true if this error is temporary and could be retried
func IsRetryable(err error) bool {
	if err == nil {
		return true
	}

	if strings.Contains(err.Error(), "field type conflict") {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package tutum

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
)

var customUserAgent = "go-tutum/" + version

func SetUserAgent(name string) string {
	customUserAgent = ""
	customUserAgent = name + " go-tutum/" + version
	return customUserAgent
}

func TutumCall(url string, requestType string, requestBody []byte) ([]byte, error) {
	if os.Getenv("TUTUM_REST_HOST") != "" {
		BaseUrl = os.Getenv("TUTUM_BASE_URL")
		BaseUrl = BaseUrl + "/api/v1/"
	} else {
		BaseUrl = "https://dashboard.tutum.co/api/v1/"
	}

	if !IsAuthenticated() {
		return nil, fmt.Errorf("Couldn't find any Tutum credentials in ~/.tutum or environment variables TUTUM_USER and TUTUM_APIKEY")
	}

	client := &http.Client{}
	req, err := http.NewRequest(requestType, BaseUrl+url, bytes.NewBuffer(requestBody))

	if ApiKey != "" {
		AuthHeader = fmt.Sprintf("ApiKey %s:%s", User, ApiKey)
	}

	req.Header.Add("Authorization", AuthHeader)
	req.Header.Add("Accept", "application/json")
	req.Header.Add("User-Agent", customUserAgent)

	response, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if response.StatusCode > 300 {
		return nil, fmt.Errorf("Failed API call: %s ", response.Status)
	}

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return data, nil
}
<commit_msg>fix typo in http call function<commit_after>package tutum

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
)

var customUserAgent = "go-tutum/" + version

func SetUserAgent(name string) string {
	customUserAgent = ""
	customUserAgent = name + " go-tutum/" + version
	return customUserAgent
}

func TutumCall(url string, requestType string, requestBody []byte) ([]byte, error) {
	if os.Getenv("TUTUM_REST_HOST") != "" {
		BaseUrl = os.Getenv("TUTUM_REST_HOST")
		BaseUrl = BaseUrl + "/api/v1/"
	} else {
		BaseUrl = "https://dashboard.tutum.co/api/v1/"
	}

	if !IsAuthenticated() {
		return nil, fmt.Errorf("Couldn't find any Tutum credentials in ~/.tutum or environment variables TUTUM_USER and TUTUM_APIKEY")
	}

	client := &http.Client{}
	req, err := http.NewRequest(requestType, BaseUrl+url, bytes.NewBuffer(requestBody))

	if ApiKey != "" {
		AuthHeader = fmt.Sprintf("ApiKey %s:%s", User, ApiKey)
	}

	req.Header.Add("Authorization", AuthHeader)
	req.Header.Add("Accept", "application/json")
	req.Header.Add("User-Agent", customUserAgent)

	response, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if response.StatusCode > 300 {
		return nil, fmt.Errorf("Failed API call: %s ", response.Status)
	}

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return data, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Unknwon
//
// 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 main

import (
	"log"
	"os"

	"github.com/Unknwon/com"
	"gopkg.in/ini.v0"
)

var cmdSync = &Command{
	UsageLine: "sync [source file] [target files]",
	Short:     "sync keys for locale files",
	Long: `to quickly sync keys for one or more locale files
based on the one you already have
`,
}

func init() {
	cmdSync.Run = syncLocales
}

func syncLocales(cmd *Command, args []string) {
	switch len(args) {
	case 0:
		log.Fatalln("No source locale file is specified")
	case 1:
		log.Fatalln("No target locale file is specified")
	}

	srcLocale, err := ini.Load(args[0])
	if err != nil {
		log.Fatalln(err)
	}

	// Load or create target locales.
	targets := args[1:]
	targetLocales := make([]*ini.File, len(targets))
	for i, target := range targets {
		if !com.IsExist(target) {
			os.Create(target)
		}

		targetLocales[i], err = ini.Load(target)
		if err != nil {
			log.Fatalln(err)
		}
	}

	for _, secName := range srcLocale.SectionStrings() {
		keyList := srcLocale.Section(secName).KeyStrings()
		for _, target := range targetLocales {
			sec, err := target.GetSection(secName)
			if err != nil {
				if sec, err = target.NewSection(secName); err != nil {
					log.Fatalln(err)
				}
			}
			for _, k := range keyList {
				if _, err = sec.GetKey(k); err != nil {
					sec.NewKey(k, "")
				}
			}
		}
	}

	for i, target := range targetLocales {
		if err = target.SaveTo(targets[i]); err != nil {
			log.Fatalln(err)
		}
	}
}
<commit_msg>Fix 	"gopkg.in/ini.v0" (#2)<commit_after>// Copyright 2013 Unknwon
//
// 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 main

import (
	"log"
	"os"

	"github.com/Unknwon/com"
	"gopkg.in/ini.v1"
)

var cmdSync = &Command{
	UsageLine: "sync [source file] [target files]",
	Short:     "sync keys for locale files",
	Long: `to quickly sync keys for one or more locale files
based on the one you already have
`,
}

func init() {
	cmdSync.Run = syncLocales
}

func syncLocales(cmd *Command, args []string) {
	switch len(args) {
	case 0:
		log.Fatalln("No source locale file is specified")
	case 1:
		log.Fatalln("No target locale file is specified")
	}

	srcLocale, err := ini.Load(args[0])
	if err != nil {
		log.Fatalln(err)
	}

	// Load or create target locales.
	targets := args[1:]
	targetLocales := make([]*ini.File, len(targets))
	for i, target := range targets {
		if !com.IsExist(target) {
			os.Create(target)
		}

		targetLocales[i], err = ini.Load(target)
		if err != nil {
			log.Fatalln(err)
		}
	}

	for _, secName := range srcLocale.SectionStrings() {
		keyList := srcLocale.Section(secName).KeyStrings()
		for _, target := range targetLocales {
			sec, err := target.GetSection(secName)
			if err != nil {
				if sec, err = target.NewSection(secName); err != nil {
					log.Fatalln(err)
				}
			}
			for _, k := range keyList {
				if _, err = sec.GetKey(k); err != nil {
					sec.NewKey(k, "")
				}
			}
		}
	}

	for i, target := range targetLocales {
		if err = target.SaveTo(targets[i]); err != nil {
			log.Fatalln(err)
		}
	}
}
<|endoftext|>


<commit_before>package apps

import (
	"fmt"
	"time"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa/events"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"
	. "github.com/cloudfoundry/cf-acceptance-tests/helpers/matchers"

	"crypto/tls"
	"strings"

	"encoding/json"
	"os"
	"path/filepath"
)

var _ = Describe("loggregator", func() {
	var appName string
	const hundredthOfOneSecond = 10000 // this app uses millionth of seconds

	BeforeEach(func() {
		appName = generator.PrefixedRandomName("CATS-APP-")

		Expect(cf.Cf("push", appName, "--no-start", "-b", config.RubyBuildpackName, "-m", DEFAULT_MEMORY_LIMIT, "-p", assets.NewAssets().LoggregatorLoadGenerator, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		app_helpers.SetBackend(appName)
		Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
	})

	AfterEach(func() {
		app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

		Expect(cf.Cf("delete", appName, "-f", "-r").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
	})

	Context("cf logs", func() {
		var logs *Session

		BeforeEach(func() {
			logs = cf.Cf("logs", appName)
		})

		AfterEach(func() {
			// logs might be nil if the BeforeEach panics
			if logs != nil {
				logs.Interrupt().Wait(DEFAULT_TIMEOUT)
			}
		})

		It("exercises basic loggregator behavior", func() {
			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Connected, tailing logs for app"))

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Muahaha"))
		})
	})

	Context("cf logs --recent", func() {
		It("makes loggregator buffer and dump log messages", func() {
			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(func() *Session {
				appLogsSession := cf.Cf("logs", "--recent", appName)
				Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				return appLogsSession
			}, DEFAULT_TIMEOUT).Should(Say("Muahaha"))
		})
	})

	Context("firehose data", func() {
		It("shows logs and metrics", func() {
			config := helpers.LoadConfig()

			noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
			msgChan := make(chan *events.Envelope, 100000)
			errorChan := make(chan error)
			stopchan := make(chan struct{})
			go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
			defer close(stopchan)

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			timeout := time.After(5 * time.Second)
			messages := make([]*events.Envelope, 0, 100000)

			for {
				select {
				case <-timeout:
					return
				case msg := <-msgChan:
					messages = append(messages, msg)
				}
			}

			Expect(messages).To(ContainElement(EnvelopeContainingMessageLike("Muahaha")), "To enable the logging & metrics firehose feature, please ask your CF administrator to add the 'doppler.firehose' scope to your CF admin user.")
		})
	})
})

type cfHomeConfig struct {
	AccessToken         string
	LoggregatorEndpoint string
}

func getCfHomeConfig() *cfHomeConfig {
	myCfHomeConfig := &cfHomeConfig{}

	cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
		path := filepath.Join(os.Getenv("CF_HOME"), ".cf", "config.json")

		configFile, err := os.Open(path)
		if err != nil {
			panic(err)
		}

		decoder := json.NewDecoder(configFile)
		err = decoder.Decode(myCfHomeConfig)
		if err != nil {
			panic(err)
		}
	})

	return myCfHomeConfig
}

func getAdminUserAccessToken() string {
	return getCfHomeConfig().AccessToken
}

func getDopplerEndpoint() string {
	return strings.Replace(getCfHomeConfig().LoggregatorEndpoint, "loggregator", "doppler", -1)
}
<commit_msg>Add test to verify container metrics<commit_after>package apps

import (
	"fmt"
	"time"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa/events"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"
	. "github.com/cloudfoundry/cf-acceptance-tests/helpers/matchers"

	"crypto/tls"
	"strings"

	"encoding/json"
	"os"
	"path/filepath"
)

var _ = Describe("loggregator", func() {
	var appName string
	const hundredthOfOneSecond = 10000 // this app uses millionth of seconds

	BeforeEach(func() {
		appName = generator.PrefixedRandomName("CATS-APP-")

		Expect(cf.Cf("push", appName, "--no-start", "-b", config.RubyBuildpackName, "-m", DEFAULT_MEMORY_LIMIT, "-p", assets.NewAssets().LoggregatorLoadGenerator, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		app_helpers.SetBackend(appName)
		Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
	})

	AfterEach(func() {
		app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

		Expect(cf.Cf("delete", appName, "-f", "-r").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
	})

	Context("cf logs", func() {
		var logs *Session

		BeforeEach(func() {
			logs = cf.Cf("logs", appName)
		})

		AfterEach(func() {
			// logs might be nil if the BeforeEach panics
			if logs != nil {
				logs.Interrupt().Wait(DEFAULT_TIMEOUT)
			}
		})

		It("exercises basic loggregator behavior", func() {
			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Connected, tailing logs for app"))

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Muahaha"))
		})
	})

	Context("cf logs --recent", func() {
		It("makes loggregator buffer and dump log messages", func() {
			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(func() *Session {
				appLogsSession := cf.Cf("logs", "--recent", appName)
				Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				return appLogsSession
			}, DEFAULT_TIMEOUT).Should(Say("Muahaha"))
		})
	})

	Context("firehose data", func() {
		It("shows logs and metrics", func() {
			config := helpers.LoadConfig()

			noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
			msgChan := make(chan *events.Envelope, 100000)
			errorChan := make(chan error)
			stopchan := make(chan struct{})
			go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
			defer close(stopchan)

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			timeout := time.After(5 * time.Second)
			messages := make([]*events.Envelope, 0, 100000)

			for {
				select {
				case <-timeout:
					return
				case msg := <-msgChan:
					messages = append(messages, msg)
				}
			}

			Expect(messages).To(ContainElement(EnvelopeContainingMessageLike("Muahaha")), "To enable the logging & metrics firehose feature, please ask your CF administrator to add the 'doppler.firehose' scope to your CF admin user.")
		})
	})

	It("shows container metrics", func() {
		config := helpers.LoadConfig()
		appGuid := strings.TrimSpace(string(cf.Cf("app", appName, "--guid").Wait(DEFAULT_TIMEOUT).Out.Contents()))

		noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
		msgChan := make(chan *events.Envelope, 100000)
		errorChan := make(chan error)
		stopchan := make(chan struct{})
		go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
		defer close(stopchan)

		Eventually(func() bool {
			for {
				select {
				case msg := <-msgChan:
					if cm := msg.GetContainerMetric(); cm != nil {
						if cm.GetApplicationId() == appGuid && cm.GetInstanceIndex() == 0 {
							return true
						}
					}
				default:
					return false
				}
			}
		}, DEFAULT_TIMEOUT).Should(BeTrue())
	})
})

type cfHomeConfig struct {
	AccessToken         string
	LoggregatorEndpoint string
}

func getCfHomeConfig() *cfHomeConfig {
	myCfHomeConfig := &cfHomeConfig{}

	cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
		path := filepath.Join(os.Getenv("CF_HOME"), ".cf", "config.json")

		configFile, err := os.Open(path)
		if err != nil {
			panic(err)
		}

		decoder := json.NewDecoder(configFile)
		err = decoder.Decode(myCfHomeConfig)
		if err != nil {
			panic(err)
		}
	})

	return myCfHomeConfig
}

func getAdminUserAccessToken() string {
	return getCfHomeConfig().AccessToken
}

func getDopplerEndpoint() string {
	return strings.Replace(getCfHomeConfig().LoggregatorEndpoint, "loggregator", "doppler", -1)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/los"
	"github.com/phil-mansfield/gotetra/los/geom"
	"github.com/phil-mansfield/gotetra/los/analyze"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/render/halo"
)

type Params struct {
	// HaloProfiles params
	RBins, Spokes, Rings int
	MaxMult, MinMult float64

	// Splashback params
	Order, Window, Levels int
	Cutoff float64
}

func main() {
	// Parse.
	p := parseCmd()
	ids, snaps, err := parseStdin()
	if err != nil { log.Fatal(err.Error()) }

	if len(ids) == 0 { return }

	// Compute coefficients.
	coeffs := make([][]float64, len(ids))
	snapBins, idxBins := binBySnap(snaps, ids)
	buf := make([]analyze.RingBuffer, p.Rings)
	for i := range buf { buf[i].Init(p.Spokes, p.RBins) }

	var losBuf *los.Buffers
	for snap, snapIDs := range snapBins {
		idxs := idxBins[snap]

		if snap == -1 { continue }

		// Bin halos
		hds, files, err := readHeaders(snap)
		if err != nil { err.Error() }
		if losBuf == nil { losBuf = los.NewBuffers(files[0], &hds[0]) }
		halos, err := createHalos(snap, &hds[0], snapIDs, p)
		if err != nil { log.Fatal(err.Error()) }
		intrBins := binIntersections(hds, halos)

		// Add densities. Done header by header to limit I/O time.
		for i := range hds {
			hdContainer := []io.SheetHeader{hds[i]}
			fileContainer := []string{files[i]}
			los.LoadPtrDensities(
				intrBins[i], hdContainer, fileContainer, losBuf,
			)
		}

		// Calculate Penna coefficients.
		for i := range halos {
			coeffs[idxs[i]] = calcCoeffs(&halos[i], buf, p)
		}
	}

	printCoeffs(ids, snaps, coeffs)
}

func parseCmd() *Params {
	// Parse command line.
	p := &Params{}
	flag.IntVar(&p.RBins, "RBins", 256,
		"Number of radial bins used per LoS.")
	flag.IntVar(&p.Spokes, "Spokes", 1024,
		"Number of LoS's used per ring.")
	flag.IntVar(&p.Rings, "Rings", 10,
		"Number of rings used per halo. 3, 4, 6, and 10 rings are\n" +
			"guaranteed to be uniformly spaced.")
	flag.Float64Var(&p.MaxMult, "MaxMult", 3,
		"Ending radius of LoSs as a multiple of R_200m.")
	flag.Float64Var(&p.MinMult, "MinMult", 0.5,
		"Starting radius of LoSs as a multiple of R_200m.")
	flag.IntVar(&p.Order, "Order", 5,
		"Order of the shell fitting function.")
	flag.IntVar(&p.Window, "Window", 121,
		"Number of bins within smoothign window. Must be odd.")
	flag.IntVar(&p.Levels, "Levels", 4,
		"The number of recurve max-finding levels used by the 2D edge finder.")
	flag.Float64Var(&p.Cutoff, "Cutoff", 0.0,
		"The shallowest slope that can be considered a splashback point.")
	flag.Parse()
	return p
}

func parseStdin() (ids, snaps []int, err error) {
	ids, snaps = []int{}, []int{}
	lines, err := stdinLines()
	if err != nil { return nil, nil, err }
	for i, line := range lines {
		tokens := strings.Split(line, " ")

		var (
			id, snap int
			err error
		)
		switch len(tokens) {
		case 0:
			continue
		case 2:
			id, err = strconv.Atoi(tokens[0])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[0],
				)
			}
			snap, err = strconv.Atoi(tokens[1])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[1],
				)
			}
		case 1:
			if tokens[0] == "" { continue }
			fallthrough
		default:
			return nil, nil, fmt.Errorf(
				"Line %d of stdin has %d tokens, but 2 are required.",
				i + 1, len(tokens),
			)
		}

		ids = append(ids, id)
		snaps = append(snaps, snap)
	}

	return ids, snaps, err
}

func stdinLines() ([]string, error) {
		bs, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, fmt.Errorf(
			"Error reading stdin: %s.", err.Error(),
		)
	}

	text := string(bs)
	return strings.Split(text, "\n"), nil
}

func readHeaders(snap int) ([]io.SheetHeader, []string, error) {
	memoDir := os.Getenv("GTET_MEMO_DIR")
	if memoDir == "" {
		// You don't want to memoize? Fine. Deal with the consequences.
		return readHeadersFromSheet(snap)
	}
	if _, err := os.Stat(memoDir); err != nil {
		return nil, nil, err
	}

	memoFile := path.Join(memoDir, fmt.Sprintf("hd_snap%d.dat", snap))

	if _, err := os.Stat(memoFile); err != nil {
		// File not written yet.
		hds, files, err := readHeadersFromSheet(snap)
		if err != nil { return nil, nil, err }

        f, err := os.Create(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()
        binary.Write(f, binary.LittleEndian, hds)

		return hds, files, nil
	} else {
		// File exists: read from it instead.

		f, err := os.Open(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()

		n, err := sheetNum(snap)
		if err != nil { return nil, nil, err }
		hds := make([]io.SheetHeader, n)
        binary.Read(f, binary.LittleEndian, hds)

		gtetFmt := os.Getenv("GTET_FMT")
		dir := fmt.Sprintf(gtetFmt, snap)
		files, err := dirContents(dir)
		if err != nil { return nil, nil, err }

		return hds, files, nil
	}

}

func dirContents(dir string) ([]string, error) {
	infos, err := ioutil.ReadDir(dir)
	if err != nil { return nil, err }

	files := make([]string, len(infos))
	for i := range infos {
		files[i] = path.Join(dir, infos[i].Name())
	}

	return files, nil
}

func sheetNum(snap int) (int, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return 0, err }
	return len(files), nil
}

func readHeadersFromSheet(snap int) ([]io.SheetHeader, []string, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return nil, nil, err }

	hds := make([]io.SheetHeader, len(files))
	for i := range files {
		err = io.ReadSheetHeaderAt(files[i], &hds[i])
		if err != nil { return nil, nil, err }
	}
	return hds, files, nil
}

func createHalos(
	snap int, hd *io.SheetHeader, ids []int, p *Params,
) ([]los.HaloProfiles, error) {
	// Read coordinates, radii and IDs of halos.
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	if rockstarDir == "" {
		return nil, fmt.Errorf("$GTET_ROCKSTAR_DIR not set.")
	}

	files, err := dirContents(rockstarDir)
	if err != nil { return nil, err }
	file := files[snap]

	ids, vals, err := halo.ReadRockstarVals(
		file, &hd.Cosmo, halo.X, halo.Y, halo.Z, halo.M200b,
	)
	if err != nil { return nil, err }

	xs, ys, zs, ms := vals[1], vals[2], vals[3], vals[4]
	halo.R200m.Radius(&hd.Cosmo, ms, ms)
	rs := ms

	// Initialize halos.
	halos := make([]los.HaloProfiles, len(ids))
	seenIDs := make(map[int]bool)
	for i, id := range ids {
		idx := sort.SearchInts(ids, id)
		if ids[idx] !=  id {
			return nil, fmt.Errorf("Halo ID %d not in halo catalog.", id)
		}

		origin := &geom.Vec{
			float32(xs[idx]), float32(ys[idx]), float32(zs[idx]),
		}

		// If we've already seen a halo once, randomize its orientation.
		if seenIDs[id] {
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
				los.Rotate(float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64())),
			)
		} else {
			seenIDs[id] = true
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
			)
		}
	}

	panic("NYI")
}

func binIntersections(
	hds []io.SheetHeader, halos []los.HaloProfiles,
) [][]*los.HaloProfiles {

	bins := make([][]*los.HaloProfiles, len(hds))
	for i := range hds {
		for hi := range halos {
			if (&halos[hi]).SheetIntersect(&hds[i]) {
				bins[i] = append(bins[i], &halos[hi])
			}
		}
	}
	return bins
}

func calcCoeffs(
	halo *los.HaloProfiles, buf []analyze.RingBuffer, p *Params,
) []float64 {
	for i := range buf {
		buf[i].Clear()
		buf[i].Splashback(halo, i, p.Window, p.Cutoff)
	}
	pxs, pys := analyze.FilterPoints(buf, p.Levels)
	cs, _ := analyze.PennaVolumeFit(pxs, pys, halo, p.Order, p.Order)
	return cs
}

func printCoeffs(ids, snaps []int, coeffs [][]float64) {
	idWidth, snapWidth := 0, 0
	coeffWidths := make([]int, len(coeffs[0]))
	for i := range ids {
		iWidth := len(fmt.Sprintf("%d", ids[i]))
		sWidth := len(fmt.Sprintf("%d", snaps[i]))
		if iWidth > idWidth { idWidth = iWidth }
		if sWidth > snapWidth { snapWidth = sWidth }
		for j := range coeffs[i] {
			width := len(fmt.Sprintf("%.5g", coeffs[i][j]))
			if width > coeffWidths[j] { coeffWidths[j] = width }
		}
	}

	rowFmt := fmt.Sprintf("%%%dd %%%dd", idWidth, snapWidth)
	coefFmts := make([]string, len(coeffs[0]))
	for i := range coefFmts {
		rowFmt += fmt.Sprintf(" %%%d.5g", coeffWidths[i])
	}
	rowFmt += "\n"
	for i := range ids {
		args := append([]interface{}{ids[i], snaps[i]}, intr(coeffs[i])...)
		fmt.Printf(rowFmt, args...)
	}
}

func intr(xs []float64) []interface{} {
	is := make([]interface{}, len(xs))
	for i, x := range xs { is[i] = x }
	return is
}

func binBySnap(snaps, ids []int) (snapBins, idxBins map[int][]int) {
	snapBins = make(map[int][]int)
	idxBins = make(map[int][]int)
	for i, snap := range snaps {
		id := ids[i]
		snapBins[snap] = append(snapBins[snap], id)
		idxBins[snap] = append(idxBins[snap], i)
	}
	return snapBins, idxBins
}
<commit_msg>I guess I made a change to this file?<commit_after>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/los"
	"github.com/phil-mansfield/gotetra/los/geom"
	"github.com/phil-mansfield/gotetra/los/analyze"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/render/halo"
)

type Params struct {
	// HaloProfiles params
	RBins, Spokes, Rings int
	MaxMult, MinMult float64

	// Splashback params
	Order, Window, Levels int
	Cutoff float64
}

func main() {
	// Parse.
	p := parseCmd()
	ids, snaps, err := parseStdin()
	if err != nil { log.Fatal(err.Error()) }

	if len(ids) == 0 { return }

	// Compute coefficients.
	coeffs := make([][]float64, len(ids))
	snapBins, idxBins := binBySnap(snaps, ids)
	buf := make([]analyze.RingBuffer, p.Rings)
	for i := range buf { buf[i].Init(p.Spokes, p.RBins) }

	var losBuf *los.Buffers
	for snap, snapIDs := range snapBins {
		idxs := idxBins[snap]

		if snap == -1 { continue }

		// Bin halos
		hds, files, err := readHeaders(snap)
		if err != nil { err.Error() }
		if losBuf == nil { losBuf = los.NewBuffers(files[0], &hds[0]) }
		halos, err := createHalos(snap, &hds[0], snapIDs, p)
		if err != nil { log.Fatal(err.Error()) }
		intrBins := binIntersections(hds, halos)

		// Add densities. Done header by header to limit I/O time.
		for i := range hds {
			hdContainer := []io.SheetHeader{hds[i]}
			fileContainer := []string{files[i]}
			los.LoadPtrDensities(
				intrBins[i], hdContainer, fileContainer, losBuf,
			)
		}

		// Calculate Penna coefficients.
		for i := range halos {
			coeffs[idxs[i]] = calcCoeffs(&halos[i], buf, p)
		}
	}

	printCoeffs(ids, snaps, coeffs)
}

func parseCmd() *Params {
	// Parse command line.
	p := &Params{}
	flag.IntVar(&p.RBins, "RBins", 256,
		"Number of radial bins used per LoS.")
	flag.IntVar(&p.Spokes, "Spokes", 1024,
		"Number of LoS's used per ring.")
	flag.IntVar(&p.Rings, "Rings", 10,
		"Number of rings used per halo. 3, 4, 6, and 10 rings are\n" +
			"guaranteed to be uniformly spaced.")
	flag.Float64Var(&p.MaxMult, "MaxMult", 3,
		"Ending radius of LoSs as a multiple of R_200m.")
	flag.Float64Var(&p.MinMult, "MinMult", 0.5,
		"Starting radius of LoSs as a multiple of R_200m.")
	flag.IntVar(&p.Order, "Order", 5,
		"Order of the shell fitting function.")
	flag.IntVar(&p.Window, "Window", 121,
		"Number of bins within smoothign window. Must be odd.")
	flag.IntVar(&p.Levels, "Levels", 4,
		"The number of recurve max-finding levels used by the 2D edge finder.")
	flag.Float64Var(&p.Cutoff, "Cutoff", 0.0,
		"The shallowest slope that can be considered a splashback point.")
	flag.Parse()
	return p
}

func parseStdin() (ids, snaps []int, err error) {
	ids, snaps = []int{}, []int{}
	lines, err := stdinLines()
	if err != nil { return nil, nil, err }
	for i, line := range lines {
		tokens := strings.Split(line, " ")

		var (
			id, snap int
			err error
		)
		switch len(tokens) {
		case 0:
			continue
		case 2:
			id, err = strconv.Atoi(tokens[0])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[0],
				)
			}
			snap, err = strconv.Atoi(tokens[1])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[1],
				)
			}
		case 1:
			if tokens[0] == "" { continue }
			fallthrough
		default:
			return nil, nil, fmt.Errorf(
				"Line %d of stdin has %d tokens, but 2 are required.",
				i + 1, len(tokens),
			)
		}

		ids = append(ids, id)
		snaps = append(snaps, snap)
	}

	return ids, snaps, nil
}

func stdinLines() ([]string, error) {
		bs, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, fmt.Errorf(
			"Error reading stdin: %s.", err.Error(),
		)
	}

	text := string(bs)
	return strings.Split(text, "\n"), nil
}

func readHeaders(snap int) ([]io.SheetHeader, []string, error) {
	memoDir := os.Getenv("GTET_MEMO_DIR")
	if memoDir == "" {
		// You don't want to memoize? Fine. Deal with the consequences.
		return readHeadersFromSheet(snap)
	}
	if _, err := os.Stat(memoDir); err != nil {
		return nil, nil, err
	}

	memoFile := path.Join(memoDir, fmt.Sprintf("hd_snap%d.dat", snap))

	if _, err := os.Stat(memoFile); err != nil {
		// File not written yet.
		hds, files, err := readHeadersFromSheet(snap)
		if err != nil { return nil, nil, err }

        f, err := os.Create(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()
        binary.Write(f, binary.LittleEndian, hds)

		return hds, files, nil
	} else {
		// File exists: read from it instead.

		f, err := os.Open(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()

		n, err := sheetNum(snap)
		if err != nil { return nil, nil, err }
		hds := make([]io.SheetHeader, n)
        binary.Read(f, binary.LittleEndian, hds)

		gtetFmt := os.Getenv("GTET_FMT")
		dir := fmt.Sprintf(gtetFmt, snap)
		files, err := dirContents(dir)
		if err != nil { return nil, nil, err }

		return hds, files, nil
	}

}

func sheetNum(snap int) (int, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return 0, err }
	return len(files), nil
}

func readHeadersFromSheet(snap int) ([]io.SheetHeader, []string, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return nil, nil, err }

	hds := make([]io.SheetHeader, len(files))
	for i := range files {
		err = io.ReadSheetHeaderAt(files[i], &hds[i])
		if err != nil { return nil, nil, err }
	}
	return hds, files, nil
}

func createHalos(
	snap int, hd *io.SheetHeader, ids []int, p *Params,
) ([]los.HaloProfiles, error) {
	// Read coordinates, radii and IDs of halos.
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	if rockstarDir == "" {
		return nil, fmt.Errorf("$GTET_ROCKSTAR_DIR not set.")
	}

	files, err := dirContents(rockstarDir)
	if err != nil { return nil, err }
	file := files[snap]

	ids, vals, err := halo.ReadRockstarVals(
		file, &hd.Cosmo, halo.X, halo.Y, halo.Z, halo.M200b,
	)
	if err != nil { return nil, err }

	xs, ys, zs, ms := vals[1], vals[2], vals[3], vals[4]
	halo.R200m.Radius(&hd.Cosmo, ms, ms)
	rs := ms

	// Initialize halos.
	halos := make([]los.HaloProfiles, len(ids))
	seenIDs := make(map[int]bool)
	for i, id := range ids {
		idx := sort.SearchInts(ids, id)
		if ids[idx] !=  id {
			return nil, fmt.Errorf("Halo ID %d not in halo catalog.", id)
		}

		origin := &geom.Vec{
			float32(xs[idx]), float32(ys[idx]), float32(zs[idx]),
		}

		// If we've already seen a halo once, randomize its orientation.
		if seenIDs[id] {
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
				los.Rotate(float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64())),
			)
		} else {
			seenIDs[id] = true
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
			)
		}
	}

	panic("NYI")
}

func binIntersections(
	hds []io.SheetHeader, halos []los.HaloProfiles,
) [][]*los.HaloProfiles {

	bins := make([][]*los.HaloProfiles, len(hds))
	for i := range hds {
		for hi := range halos {
			if (&halos[hi]).SheetIntersect(&hds[i]) {
				bins[i] = append(bins[i], &halos[hi])
			}
		}
	}
	return bins
}

func calcCoeffs(
	halo *los.HaloProfiles, buf []analyze.RingBuffer, p *Params,
) []float64 {
	for i := range buf {
		buf[i].Clear()
		buf[i].Splashback(halo, i, p.Window, p.Cutoff)
	}
	pxs, pys := analyze.FilterPoints(buf, p.Levels)
	cs, _ := analyze.PennaVolumeFit(pxs, pys, halo, p.Order, p.Order)
	return cs
}

func printCoeffs(ids, snaps []int, coeffs [][]float64) {
	idWidth, snapWidth := 0, 0
	coeffWidths := make([]int, len(coeffs[0]))
	for i := range ids {
		iWidth := len(fmt.Sprintf("%d", ids[i]))
		sWidth := len(fmt.Sprintf("%d", snaps[i]))
		if iWidth > idWidth { idWidth = iWidth }
		if sWidth > snapWidth { snapWidth = sWidth }
		for j := range coeffs[i] {
			width := len(fmt.Sprintf("%.5g", coeffs[i][j]))
			if width > coeffWidths[j] { coeffWidths[j] = width }
		}
	}

	rowFmt := fmt.Sprintf("%%%dd %%%dd", idWidth, snapWidth)
	coefFmts := make([]string, len(coeffs[0]))
	for i := range coefFmts {
		rowFmt += fmt.Sprintf(" %%%d.5g", coeffWidths[i])
	}
	rowFmt += "\n"
	for i := range ids {
		args := append([]interface{}{ids[i], snaps[i]}, intr(coeffs[i])...)
		fmt.Printf(rowFmt, args...)
	}
}

func intr(xs []float64) []interface{} {
	is := make([]interface{}, len(xs))
	for i, x := range xs { is[i] = x }
	return is
}

func binBySnap(snaps, ids []int) (snapBins, idxBins map[int][]int) {
	snapBins = make(map[int][]int)
	idxBins = make(map[int][]int)
	for i, snap := range snaps {
		id := ids[i]
		snapBins[snap] = append(snapBins[snap], id)
		idxBins[snap] = append(idxBins[snap], i)
	}
	return snapBins, idxBins
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/mux"

	"github.com/lxc/lxd/shared"
)

type containerStatePutReq struct {
	Action   string `json:"action"`
	Timeout  int    `json:"timeout"`
	Force    bool   `json:"force"`
	Stateful bool   `json:"stateful"`
}

func containerState(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]
	c, err := containerLoadByName(d, name)
	if err != nil {
		return SmartError(err)
	}

	state, err := c.RenderState()
	if err != nil {
		return InternalError(err)
	}

	return SyncResponse(true, state)
}

func containerStatePut(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	raw := containerStatePutReq{}

	// We default to -1 (i.e. no timeout) here instead of 0 (instant
	// timeout).
	raw.Timeout = -1

	if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
		return BadRequest(err)
	}

	c, err := containerLoadByName(d, name)
	if err != nil {
		return SmartError(err)
	}

	var do func(*operation) error
	switch shared.ContainerAction(raw.Action) {
	case shared.Start:
		do = func(op *operation) error {
			if err = c.Start(raw.Stateful); err != nil {
				return err
			}
			return nil
		}
	case shared.Stop:
		if raw.Stateful {
			do = func(op *operation) error {
				err := c.Stop(raw.Stateful)
				if err != nil {
					return err
				}

				return nil
			}
		} else if raw.Timeout == 0 || raw.Force {
			do = func(op *operation) error {
				err = c.Stop(false)
				if err != nil {
					return err
				}

				if c.IsEphemeral() {
					c.Delete()
				}

				return nil
			}
		} else {
			do = func(op *operation) error {
				err = c.Shutdown(time.Duration(raw.Timeout) * time.Second)
				if err != nil {
					return err
				}

				if c.IsEphemeral() {
					c.Delete()
				}

				return nil
			}
		}
	case shared.Restart:
		do = func(op *operation) error {
			if raw.Timeout == 0 || raw.Force {
				err = c.Stop(false)
				if err != nil {
					return err
				}
			} else {
				err = c.Shutdown(time.Duration(raw.Timeout) * time.Second)
				if err != nil {
					return err
				}
			}
			err = c.Start(false)
			if err != nil {
				return err
			}

			return nil
		}
	case shared.Freeze:
		do = func(op *operation) error {
			return c.Freeze()
		}
	case shared.Unfreeze:
		do = func(op *operation) error {
			return c.Unfreeze()
		}
	default:
		return BadRequest(fmt.Errorf("unknown action %s", raw.Action))
	}

	resources := map[string][]string{}
	resources["containers"] = []string{name}

	op, err := operationCreate(operationClassTask, resources, nil, do, nil, nil)
	if err != nil {
		return InternalError(err)
	}

	return OperationResponse(op)
}
<commit_msg>Prevent container actions while in setup mode<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/mux"

	"github.com/lxc/lxd/shared"
)

type containerStatePutReq struct {
	Action   string `json:"action"`
	Timeout  int    `json:"timeout"`
	Force    bool   `json:"force"`
	Stateful bool   `json:"stateful"`
}

func containerState(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]
	c, err := containerLoadByName(d, name)
	if err != nil {
		return SmartError(err)
	}

	state, err := c.RenderState()
	if err != nil {
		return InternalError(err)
	}

	return SyncResponse(true, state)
}

func containerStatePut(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	raw := containerStatePutReq{}

	// We default to -1 (i.e. no timeout) here instead of 0 (instant
	// timeout).
	raw.Timeout = -1

	if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
		return BadRequest(err)
	}

	// Don't mess with containers while in setup mode
	<-d.readyChan

	c, err := containerLoadByName(d, name)
	if err != nil {
		return SmartError(err)
	}

	var do func(*operation) error
	switch shared.ContainerAction(raw.Action) {
	case shared.Start:
		do = func(op *operation) error {
			if err = c.Start(raw.Stateful); err != nil {
				return err
			}
			return nil
		}
	case shared.Stop:
		if raw.Stateful {
			do = func(op *operation) error {
				err := c.Stop(raw.Stateful)
				if err != nil {
					return err
				}

				return nil
			}
		} else if raw.Timeout == 0 || raw.Force {
			do = func(op *operation) error {
				err = c.Stop(false)
				if err != nil {
					return err
				}

				if c.IsEphemeral() {
					c.Delete()
				}

				return nil
			}
		} else {
			do = func(op *operation) error {
				err = c.Shutdown(time.Duration(raw.Timeout) * time.Second)
				if err != nil {
					return err
				}

				if c.IsEphemeral() {
					c.Delete()
				}

				return nil
			}
		}
	case shared.Restart:
		do = func(op *operation) error {
			if raw.Timeout == 0 || raw.Force {
				err = c.Stop(false)
				if err != nil {
					return err
				}
			} else {
				err = c.Shutdown(time.Duration(raw.Timeout) * time.Second)
				if err != nil {
					return err
				}
			}
			err = c.Start(false)
			if err != nil {
				return err
			}

			return nil
		}
	case shared.Freeze:
		do = func(op *operation) error {
			return c.Freeze()
		}
	case shared.Unfreeze:
		do = func(op *operation) error {
			return c.Unfreeze()
		}
	default:
		return BadRequest(fmt.Errorf("unknown action %s", raw.Action))
	}

	resources := map[string][]string{}
	resources["containers"] = []string{name}

	op, err := operationCreate(operationClassTask, resources, nil, do, nil, nil)
	if err != nil {
		return InternalError(err)
	}

	return OperationResponse(op)
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You under the Apache License, Version 2.0
(the "License"); you may not use this file except in compliance with
the License.  You may obtain a copy of the License at

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

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

package go_kafka_client

import (
	"testing"
)

var (
	consumers         = []string{"consumerid1", "consumerid2"}
	consumerThreadIds = []ConsumerThreadId{
		ConsumerThreadId{"consumerid1", 0},
		ConsumerThreadId{"consumerid1", 1},
		ConsumerThreadId{"consumerid2", 0},
		ConsumerThreadId{"consumerid2", 1},
	}
	partitionsForTopic = map[string][]int32{
		"topic1": []int32{0, 1, 2, 3, 4},
	}
	consumersForTopic = map[string][]ConsumerThreadId{
		"topic1": consumerThreadIds,
	}
	totalPartitions = 5
)

func TestRoundRobinAssignor(t *testing.T) {
	//basic scenario
	assignor := newPartitionAssignor("roundrobin")
	context := &assignmentContext{
		Group:              "group",
		PartitionsForTopic: partitionsForTopic,
		ConsumersForTopic:  consumersForTopic,
		Consumers:          consumers,
	}

	var totalDecisions = 0
	for _, consumer := range consumers {
		context.ConsumerId = consumer
		context.MyTopicThreadIds = map[string][]ConsumerThreadId{
			"topic1": []ConsumerThreadId{
				ConsumerThreadId{consumer, 0},
				ConsumerThreadId{consumer, 1}},
		}
		ownershipDecision := assignor(context)
		decisionsNum := len(ownershipDecision)
		if decisionsNum == totalPartitions {
			t.Errorf("Too many partitions assigned to consumer %s", consumer)
		}

		totalDecisions += decisionsNum
		t.Logf("%v\n", ownershipDecision)
	}

	assert(t, totalDecisions, totalPartitions)

	//NOT every topic has the same number of streams within a consumer instance
	failed := false
	defer func() {
		r := recover()
		if r != nil {
			failed = true
		}
	}()

	context.ConsumerId = "consumerid1"
	context.MyTopicThreadIds = map[string][]ConsumerThreadId{
		"topic1": []ConsumerThreadId{
			ConsumerThreadId{"consumerid1", 0},
			ConsumerThreadId{"consumerid1", 1}},
	}
	context.ConsumersForTopic = map[string][]ConsumerThreadId{
		"topic1": consumerThreadIds,
		"topic2": []ConsumerThreadId{
			ConsumerThreadId{"consumerid1", 0},
			ConsumerThreadId{"consumerid1", 1},
			ConsumerThreadId{"consumerid2", 0},
		},
	}
	assignor(context)

	assert(t, failed, true)
}

func TestRangeAssignor(t *testing.T) {
	//basic scenario
	assignor := newPartitionAssignor("range")
	context := &assignmentContext{
		Group:              "group",
		PartitionsForTopic: partitionsForTopic,
		ConsumersForTopic:  consumersForTopic,
		Consumers:          consumers,
	}

	var totalDecisions = 0
	for _, consumer := range consumers {
		context.ConsumerId = consumer
		context.MyTopicThreadIds = map[string][]ConsumerThreadId{
			"topic1": []ConsumerThreadId{
				ConsumerThreadId{consumer, 0},
				ConsumerThreadId{consumer, 1}},
		}
		ownershipDecision := assignor(context)
		decisionsNum := len(ownershipDecision)
		if decisionsNum == totalPartitions {
			t.Errorf("too many partitions assigned to consumer %s", consumer)
		}

		totalDecisions += decisionsNum
		t.Logf("%v\n", ownershipDecision)
	}

	assert(t, totalDecisions, totalPartitions)
}
<commit_msg>added negative test for git issue #62. re #4<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You under the Apache License, Version 2.0
(the "License"); you may not use this file except in compliance with
the License.  You may obtain a copy of the License at

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

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

package go_kafka_client

import (
	"testing"
)

var (
	consumers         = []string{"consumerid1", "consumerid2"}
	consumerThreadIds = []ConsumerThreadId{
		ConsumerThreadId{"consumerid1", 0},
		ConsumerThreadId{"consumerid1", 1},
		ConsumerThreadId{"consumerid2", 0},
		ConsumerThreadId{"consumerid2", 1},
	}
	partitionsForTopic = map[string][]int32{
		"topic1": []int32{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
	}
	consumersForTopic = map[string][]ConsumerThreadId{
		"topic1": consumerThreadIds,
	}
	totalPartitions = 10
)

func TestRoundRobinAssignor(t *testing.T) {
	//basic scenario
	assignor := newPartitionAssignor("roundrobin")
	context := &assignmentContext{
		Group:              "group",
		PartitionsForTopic: partitionsForTopic,
		ConsumersForTopic:  consumersForTopic,
		Consumers:          consumers,
	}

    assignments := make(map[TopicAndPartition]string)
	var totalDecisions = 0
	for _, consumer := range consumers {
		context.ConsumerId = consumer
		context.MyTopicThreadIds = map[string][]ConsumerThreadId{
			"topic1": []ConsumerThreadId{
				ConsumerThreadId{consumer, 0},
				ConsumerThreadId{consumer, 1}},
		}
		ownershipDecision := assignor(context)
		decisionsNum := len(ownershipDecision)
		if decisionsNum == totalPartitions {
			t.Errorf("Too many partitions assigned to consumer %s", consumer)
		}

        for topicAndPartition := range ownershipDecision {
            if owner, exists := assignments[topicAndPartition]; exists {
                t.Errorf("Consumer %s tried to own topic %s and partition %d previously owned by consumer %s", consumer, topicAndPartition.Topic, topicAndPartition.Partition, owner)
            }
            assignments[topicAndPartition] = consumer
        }

		totalDecisions += decisionsNum
	}

	assert(t, totalDecisions, totalPartitions)

	//NOT every topic has the same number of streams within a consumer instance
	failed := false
	defer func() {
		r := recover()
		if r != nil {
			failed = true
		}
	}()

	context.ConsumerId = "consumerid1"
	context.MyTopicThreadIds = map[string][]ConsumerThreadId{
		"topic1": []ConsumerThreadId{
			ConsumerThreadId{"consumerid1", 0},
			ConsumerThreadId{"consumerid1", 1}},
	}
	context.ConsumersForTopic = map[string][]ConsumerThreadId{
		"topic1": consumerThreadIds,
		"topic2": []ConsumerThreadId{
			ConsumerThreadId{"consumerid1", 0},
			ConsumerThreadId{"consumerid1", 1},
			ConsumerThreadId{"consumerid2", 0},
		},
	}
	assignor(context)

	assert(t, failed, true)
}

func TestRangeAssignor(t *testing.T) {
	//basic scenario
	assignor := newPartitionAssignor("range")
	context := &assignmentContext{
		Group:              "group",
		PartitionsForTopic: partitionsForTopic,
		ConsumersForTopic:  consumersForTopic,
		Consumers:          consumers,
	}

	var totalDecisions = 0
	for _, consumer := range consumers {
		context.ConsumerId = consumer
		context.MyTopicThreadIds = map[string][]ConsumerThreadId{
			"topic1": []ConsumerThreadId{
				ConsumerThreadId{consumer, 0},
				ConsumerThreadId{consumer, 1}},
		}
		ownershipDecision := assignor(context)
		decisionsNum := len(ownershipDecision)
		if decisionsNum == totalPartitions {
			t.Errorf("too many partitions assigned to consumer %s", consumer)
		}

		totalDecisions += decisionsNum
		t.Logf("%v\n", ownershipDecision)
	}

	assert(t, totalDecisions, totalPartitions)
}
<|endoftext|>