<commit_before>package grohl

import (
	"io"
	"os"
	"path/filepath"
	"time"
)

type Data map[string]interface{}

type Logger interface {
	Log(Data) error
}

var CurrentLogger Logger = NewIoLogger(nil)
var CurrentContext = &Context{make(Data), CurrentLogger, "s", nil}
var CurrentStatter Statter = CurrentContext

func Log(data Data) {
	CurrentContext.Log(data)
}

func Report(err error, data Data) {
	CurrentContext.Report(err, data)
}

func Counter(sampleRate float32, bucket string, n ...int) {
	CurrentStatter.Counter(sampleRate, bucket, n...)
}

func Timing(sampleRate float32, bucket string, d ...time.Duration) {
	CurrentStatter.Timing(sampleRate, bucket, d...)
}

func Gauge(sampleRate float32, bucket string, value ...string) {
	CurrentStatter.Gauge(sampleRate, bucket, value...)
}

func SetLogger(logger Logger) Logger {
	if logger == nil {
		logger = NewIoLogger(nil)
	}

	CurrentLogger = logger
	CurrentContext.Logger = logger

	return logger
}

func NewChannelLogger(channel chan Data) (*ChannelLogger, chan Data) {
	if channel == nil {
		channel = make(chan Data)
	}
	return &ChannelLogger{channel}, channel
}

func NewIoLogger(stream io.Writer) *IoLogger {
	if stream == nil {
		stream = os.Stdout
	}

	return &IoLogger{stream, true}
}

func NewBufferedLogger(filename string, upperbound int) *BufferedLogger {
	isFile := len(filename) > 0

	if isFile {
		os.MkdirAll(filepath.Dir(filename), 0744)
	}

	return &BufferedLogger{
		Filename:   filename,
		Position:   0,
		Upperbound: upperbound,
		AddTime:    true,
		Fallback:   &nopcloser{os.Stdout},
		buffer:     make([]byte, upperbound),
		isFile:     isFile,
	}
}

func WriteBufferedLogs(ch chan Data, filename string, upperbound int) {
	logger := NewBufferedLogger(filename, upperbound)
	for {
		logger.Log(<-ch)
	}
}

func NewContext(data Data) *Context {
	return CurrentContext.New(data)
}

func AddContext(key string, value interface{}) {
	CurrentContext.Add(key, value)
}

func DeleteContext(key string) {
	CurrentContext.Delete(key)
}

func NewTimer(data Data) *Timer {
	return CurrentContext.Timer(data)
}

func SetTimeUnit(unit string) {
	CurrentContext.TimeUnit = unit
}

func TimeUnit() string {
	return CurrentContext.TimeUnit
}
<commit_msg>flush on restarts, terminates, and log reopens<commit_after>package grohl

import (
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"
)

type Data map[string]interface{}

type Logger interface {
	Log(Data) error
}

var CurrentLogger Logger = NewIoLogger(nil)
var CurrentContext = &Context{make(Data), CurrentLogger, "s", nil}
var CurrentStatter Statter = CurrentContext

func Log(data Data) {
	CurrentContext.Log(data)
}

func Report(err error, data Data) {
	CurrentContext.Report(err, data)
}

func Counter(sampleRate float32, bucket string, n ...int) {
	CurrentStatter.Counter(sampleRate, bucket, n...)
}

func Timing(sampleRate float32, bucket string, d ...time.Duration) {
	CurrentStatter.Timing(sampleRate, bucket, d...)
}

func Gauge(sampleRate float32, bucket string, value ...string) {
	CurrentStatter.Gauge(sampleRate, bucket, value...)
}

func SetLogger(logger Logger) Logger {
	if logger == nil {
		logger = NewIoLogger(nil)
	}

	CurrentLogger = logger
	CurrentContext.Logger = logger

	return logger
}

func NewChannelLogger(channel chan Data) (*ChannelLogger, chan Data) {
	if channel == nil {
		channel = make(chan Data)
	}
	return &ChannelLogger{channel}, channel
}

func NewIoLogger(stream io.Writer) *IoLogger {
	if stream == nil {
		stream = os.Stdout
	}

	return &IoLogger{stream, true}
}

func NewBufferedLogger(filename string, upperbound int) *BufferedLogger {
	isFile := len(filename) > 0

	if isFile {
		os.MkdirAll(filepath.Dir(filename), 0744)
	}

	return &BufferedLogger{
		Filename:   filename,
		Position:   0,
		Upperbound: upperbound,
		AddTime:    true,
		Fallback:   &nopcloser{os.Stdout},
		buffer:     make([]byte, upperbound),
		isFile:     isFile,
	}
}

func WriteBufferedLogs(logch chan Data, filename string, upperbound int) {
	logger := NewBufferedLogger(filename, upperbound)
	sigch := make(chan os.Signal)
	signal.Notify(sigch, syscall.SIGUSR1, syscall.SIGUSR2, syscall.SIGTERM)

	for {
		select {
		case <-sigch:
			logger.Flush()
		case data := <-logch:
			logger.Log(data)
		}
	}
}

func NewContext(data Data) *Context {
	return CurrentContext.New(data)
}

func AddContext(key string, value interface{}) {
	CurrentContext.Add(key, value)
}

func DeleteContext(key string) {
	CurrentContext.Delete(key)
}

func NewTimer(data Data) *Timer {
	return CurrentContext.Timer(data)
}

func SetTimeUnit(unit string) {
	CurrentContext.TimeUnit = unit
}

func TimeUnit() string {
	return CurrentContext.TimeUnit
}
<|endoftext|>


<commit_before>package supervisor

import (
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"path"
	"strconv"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cavaliercoder/grab"
	"github.com/parnurzeal/gorequest"
	"github.com/resin-io/edge-node-manager/config"
)

var (
	address string
	version string
	key     string
	rawKey  string
)

func WaitUntilReady() {
	log.Info("Waiting until supervisor is ready")

	for {
		resp, _, errs := gorequest.New().Timeout(1 * time.Second).Get(address).End()
		if errs == nil && resp.StatusCode == 401 {
			// The supervisor is up once a 401 status code is returned
			log.Info("Supervisor is ready")
			return
		}

		time.Sleep(1 * time.Second)
	}
}

func DependentApplicationsList() ([]byte, []error) {
	url, err := buildPath(address, []string{version, "dependent-apps"})
	if err != nil {
		return nil, []error{err}
	}

	req := gorequest.New()
	req.Get(url)
	req.Query(key)

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
	}).Debug("Requesting dependent applications list")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 200); errs != nil {
		return nil, errs
	}

	return body, nil
}

// DependentApplicationUpdate downloads the binary.tar for a specific application and target commit
// Saving it to {ENM_ASSETS_DIRECTORY}/{applicationUUID}/{targetCommit}/binary.tar
func DependentApplicationUpdate(applicationUUID int, targetCommit string) error {
	url, err := buildPath(address, []string{version, "dependent-apps", strconv.Itoa(applicationUUID), "assets", targetCommit})
	if err != nil {
		return err
	}

	req, err := grab.NewRequest(url)
	if err != nil {
		return err
	}

	q := req.HTTPRequest.URL.Query()
	q.Set("apikey", rawKey)
	req.HTTPRequest.URL.RawQuery = q.Encode()

	filePath := config.GetAssetsDir()
	filePath = path.Join(filePath, strconv.Itoa(applicationUUID))
	filePath = path.Join(filePath, targetCommit)
	if err = os.MkdirAll(filePath, os.ModePerm); err != nil {
		return err
	}
	filePath = path.Join(filePath, "binary.tar")
	req.Filename = filePath

	log.WithFields(log.Fields{
		"URL":         req.HTTPRequest.URL,
		"Method":      req.HTTPRequest.Method,
		"Query":       req.HTTPRequest.URL.RawQuery,
		"Destination": req.Filename,
	}).Debug("Requesting dependent application update")

	client := grab.NewClient()
	resp, err := client.Do(req)

	if err != nil {
		return err
	}

	if resp.HTTPResponse.StatusCode != 200 {
		return fmt.Errorf("Dependent application update failed")
	}

	log.Debug("Dependent application update succeeded")

	return nil
}

func DependentDeviceLog(UUID, message string) []error {
	url, err := buildPath(address, []string{version, "devices", UUID, "logs"})
	if err != nil {
		return []error{err}
	}

	type dependentDeviceLog struct {
		Message   string `json:"message"`
		TimeStamp int64  `json:"timestamp"`
	}

	content := &dependentDeviceLog{
		Message:   message,
		TimeStamp: time.Now().UTC().Unix(),
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		return []error{err}
	}

	req := gorequest.New()
	req.Post(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Transmitting dependent device log")

	resp, _, errs := req.End()
	return handleResp(resp, errs, 202)
}

func DependentDeviceInfoUpdateWithOnlineState(UUID, status, commit string, online bool) []error {
	url, err := buildPath(address, []string{version, "devices", UUID})
	if err != nil {
		return []error{err}
	}

	type dependentDeviceInfo struct {
		Status string `json:"status"`
		Online bool   `json:"is_online"`
		Commit string `json:"commit,omitempty"`
	}

	content := &dependentDeviceInfo{
		Status: status,
		Online: online,
		Commit: commit,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		return []error{err}
	}

	req := gorequest.New()
	req.Put(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Transmitting dependent device info")

	resp, _, errs := req.End()
	return handleResp(resp, errs, 200)
}

func DependentDeviceInfoUpdateWithoutOnlineState(UUID, status, commit string) []error {
	url, err := buildPath(address, []string{version, "devices", UUID})
	if err != nil {
		return []error{err}
	}

	type dependentDeviceInfo struct {
		Status string `json:"status"`
		Commit string `json:"commit,omitempty"`
	}

	content := &dependentDeviceInfo{
		Status: status,
		Commit: commit,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		return []error{err}
	}

	req := gorequest.New()
	req.Put(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Transmitting dependent device info")

	resp, _, errs := req.End()
	return handleResp(resp, errs, 200)
}

func DependentDeviceInfo(UUID string) ([]byte, []error) {
	url, err := buildPath(address, []string{version, "devices", UUID})
	if err != nil {
		return nil, []error{err}
	}

	req := gorequest.New()
	req.Get(url)
	req.Query(key)

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
	}).Debug("Requesting dependent device info")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 200); errs != nil {
		return nil, errs
	}

	return body, nil
}

func DependentDeviceProvision(applicationUUID int) (resinUUID, name string, errs []error) {
	url, err := buildPath(address, []string{version, "devices"})
	if err != nil {
		errs = []error{err}
		return
	}

	type dependentDeviceProvision struct {
		ApplicationUUID int `json:"appId"`
	}

	content := &dependentDeviceProvision{
		ApplicationUUID: applicationUUID,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		errs = []error{err}
		return
	}

	req := gorequest.New()
	req.Post(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Requesting dependent device provision")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 201); errs != nil {
		return
	}

	var buffer map[string]interface{}
	if err := json.Unmarshal(body, &buffer); err != nil {
		errs = []error{err}
		return
	}

	resinUUID = buffer["uuid"].(string)
	name = buffer["name"].(string)

	return
}

func DependentDevicesList() ([]byte, []error) {
	url, err := buildPath(address, []string{version, "devices"})
	if err != nil {
		return nil, []error{err}
	}

	req := gorequest.New()
	req.Get(url)
	req.Query(key)

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
	}).Debug("Requesting dependent devices list")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 200); errs != nil {
		return nil, errs
	}

	return body, nil
}

func init() {
	log.SetLevel(config.GetLogLevel())

	address = config.GetSuperAddr()
	version = config.GetVersion()
	rawKey = config.GetSuperAPIKey()

	type apiKey struct {
		APIKey string `json:"apikey"`
	}

	content := &apiKey{
		APIKey: rawKey,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		log.WithFields(log.Fields{
			"Key":   rawKey,
			"Error": err,
		}).Fatal("Unable to marshall API key")
	}
	key = (string)(bytes)

	log.WithFields(log.Fields{
		"Address": address,
		"Version": version,
		"Key":     key,
		"Raw key": rawKey,
	}).Debug("Initialised outgoing supervisor API")
}

func buildPath(base string, paths []string) (string, error) {
	url, err := url.ParseRequestURI(address)
	if err != nil {
		return "", err
	}

	for _, p := range paths {
		url.Path = path.Join(url.Path, p)
	}

	return url.String(), nil
}

func handleResp(resp gorequest.Response, errs []error, statusCode int) []error {
	if errs != nil {
		return errs
	}

	// Allow 404 and 410 here as it means the dep. app or dep. device has just been deleted
	if resp.StatusCode != statusCode && resp.StatusCode != 404 && resp.StatusCode != 410 {
		return []error{fmt.Errorf("Invalid response received: %s", resp.Status)}
	}

	log.WithFields(log.Fields{
		"Response": resp.Status,
	}).Debug("Valid response received")

	return nil
}
<commit_msg>Removed timestamp so that the supervisor will set the timestamp instead<commit_after>package supervisor

import (
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"path"
	"strconv"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cavaliercoder/grab"
	"github.com/parnurzeal/gorequest"
	"github.com/resin-io/edge-node-manager/config"
)

var (
	address string
	version string
	key     string
	rawKey  string
)

func WaitUntilReady() {
	log.Info("Waiting until supervisor is ready")

	for {
		resp, _, errs := gorequest.New().Timeout(1 * time.Second).Get(address).End()
		if errs == nil && resp.StatusCode == 401 {
			// The supervisor is up once a 401 status code is returned
			log.Info("Supervisor is ready")
			return
		}

		time.Sleep(1 * time.Second)
	}
}

func DependentApplicationsList() ([]byte, []error) {
	url, err := buildPath(address, []string{version, "dependent-apps"})
	if err != nil {
		return nil, []error{err}
	}

	req := gorequest.New()
	req.Get(url)
	req.Query(key)

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
	}).Debug("Requesting dependent applications list")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 200); errs != nil {
		return nil, errs
	}

	return body, nil
}

// DependentApplicationUpdate downloads the binary.tar for a specific application and target commit
// Saving it to {ENM_ASSETS_DIRECTORY}/{applicationUUID}/{targetCommit}/binary.tar
func DependentApplicationUpdate(applicationUUID int, targetCommit string) error {
	url, err := buildPath(address, []string{version, "dependent-apps", strconv.Itoa(applicationUUID), "assets", targetCommit})
	if err != nil {
		return err
	}

	req, err := grab.NewRequest(url)
	if err != nil {
		return err
	}

	q := req.HTTPRequest.URL.Query()
	q.Set("apikey", rawKey)
	req.HTTPRequest.URL.RawQuery = q.Encode()

	filePath := config.GetAssetsDir()
	filePath = path.Join(filePath, strconv.Itoa(applicationUUID))
	filePath = path.Join(filePath, targetCommit)
	if err = os.MkdirAll(filePath, os.ModePerm); err != nil {
		return err
	}
	filePath = path.Join(filePath, "binary.tar")
	req.Filename = filePath

	log.WithFields(log.Fields{
		"URL":         req.HTTPRequest.URL,
		"Method":      req.HTTPRequest.Method,
		"Query":       req.HTTPRequest.URL.RawQuery,
		"Destination": req.Filename,
	}).Debug("Requesting dependent application update")

	client := grab.NewClient()
	resp, err := client.Do(req)

	if err != nil {
		return err
	}

	if resp.HTTPResponse.StatusCode != 200 {
		return fmt.Errorf("Dependent application update failed")
	}

	log.Debug("Dependent application update succeeded")

	return nil
}

func DependentDeviceLog(UUID, message string) []error {
	url, err := buildPath(address, []string{version, "devices", UUID, "logs"})
	if err != nil {
		return []error{err}
	}

	type dependentDeviceLog struct {
		Message   string `json:"message"`
	}

	content := &dependentDeviceLog{
		Message:   message,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		return []error{err}
	}

	req := gorequest.New()
	req.Post(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Transmitting dependent device log")

	resp, _, errs := req.End()
	return handleResp(resp, errs, 202)
}

func DependentDeviceInfoUpdateWithOnlineState(UUID, status, commit string, online bool) []error {
	url, err := buildPath(address, []string{version, "devices", UUID})
	if err != nil {
		return []error{err}
	}

	type dependentDeviceInfo struct {
		Status string `json:"status"`
		Online bool   `json:"is_online"`
		Commit string `json:"commit,omitempty"`
	}

	content := &dependentDeviceInfo{
		Status: status,
		Online: online,
		Commit: commit,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		return []error{err}
	}

	req := gorequest.New()
	req.Put(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Transmitting dependent device info")

	resp, _, errs := req.End()
	return handleResp(resp, errs, 200)
}

func DependentDeviceInfoUpdateWithoutOnlineState(UUID, status, commit string) []error {
	url, err := buildPath(address, []string{version, "devices", UUID})
	if err != nil {
		return []error{err}
	}

	type dependentDeviceInfo struct {
		Status string `json:"status"`
		Commit string `json:"commit,omitempty"`
	}

	content := &dependentDeviceInfo{
		Status: status,
		Commit: commit,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		return []error{err}
	}

	req := gorequest.New()
	req.Put(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Transmitting dependent device info")

	resp, _, errs := req.End()
	return handleResp(resp, errs, 200)
}

func DependentDeviceInfo(UUID string) ([]byte, []error) {
	url, err := buildPath(address, []string{version, "devices", UUID})
	if err != nil {
		return nil, []error{err}
	}

	req := gorequest.New()
	req.Get(url)
	req.Query(key)

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
	}).Debug("Requesting dependent device info")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 200); errs != nil {
		return nil, errs
	}

	return body, nil
}

func DependentDeviceProvision(applicationUUID int) (resinUUID, name string, errs []error) {
	url, err := buildPath(address, []string{version, "devices"})
	if err != nil {
		errs = []error{err}
		return
	}

	type dependentDeviceProvision struct {
		ApplicationUUID int `json:"appId"`
	}

	content := &dependentDeviceProvision{
		ApplicationUUID: applicationUUID,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		errs = []error{err}
		return
	}

	req := gorequest.New()
	req.Post(url)
	req.Set("Content-Type", "application/json")
	req.Query(key)
	req.Send((string)(bytes))

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
		"Body":   (string)(bytes),
	}).Debug("Requesting dependent device provision")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 201); errs != nil {
		return
	}

	var buffer map[string]interface{}
	if err := json.Unmarshal(body, &buffer); err != nil {
		errs = []error{err}
		return
	}

	resinUUID = buffer["uuid"].(string)
	name = buffer["name"].(string)

	return
}

func DependentDevicesList() ([]byte, []error) {
	url, err := buildPath(address, []string{version, "devices"})
	if err != nil {
		return nil, []error{err}
	}

	req := gorequest.New()
	req.Get(url)
	req.Query(key)

	log.WithFields(log.Fields{
		"URL":    req.Url,
		"Method": req.Method,
		"Query":  req.QueryData,
	}).Debug("Requesting dependent devices list")

	resp, body, errs := req.EndBytes()
	if errs = handleResp(resp, errs, 200); errs != nil {
		return nil, errs
	}

	return body, nil
}

func init() {
	log.SetLevel(config.GetLogLevel())

	address = config.GetSuperAddr()
	version = config.GetVersion()
	rawKey = config.GetSuperAPIKey()

	type apiKey struct {
		APIKey string `json:"apikey"`
	}

	content := &apiKey{
		APIKey: rawKey,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		log.WithFields(log.Fields{
			"Key":   rawKey,
			"Error": err,
		}).Fatal("Unable to marshall API key")
	}
	key = (string)(bytes)

	log.WithFields(log.Fields{
		"Address": address,
		"Version": version,
		"Key":     key,
		"Raw key": rawKey,
	}).Debug("Initialised outgoing supervisor API")
}

func buildPath(base string, paths []string) (string, error) {
	url, err := url.ParseRequestURI(address)
	if err != nil {
		return "", err
	}

	for _, p := range paths {
		url.Path = path.Join(url.Path, p)
	}

	return url.String(), nil
}

func handleResp(resp gorequest.Response, errs []error, statusCode int) []error {
	if errs != nil {
		return errs
	}

	// Allow 404 and 410 here as it means the dep. app or dep. device has just been deleted
	if resp.StatusCode != statusCode && resp.StatusCode != 404 && resp.StatusCode != 410 {
		return []error{fmt.Errorf("Invalid response received: %s", resp.Status)}
	}

	log.WithFields(log.Fields{
		"Response": resp.Status,
	}).Debug("Valid response received")

	return nil
}
<|endoftext|>


<commit_before>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 (
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"
    "fmt"
    "os"
)

// Deletes all of the objects in the specified S3 Bucket in the region configured in the shared config
// or AWS_REGION environment variable.
//
// Usage:
//    go run s3_delete_objects BUCKET
func main() {
    if len(os.Args) != 2 {
        exitErrorf("Bucket name required\nUsage: %s BUCKET",
            os.Args[0])
    }

    bucket := os.Args[1]

    // Initialize a session in us-west-2 that the SDK will use to load
    // credentials from the shared credentials file ~/.aws/credentials.
    sess, _ := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )

    // Create S3 service client
    svc := s3.New(sess)

    // Get the list of objects
    // Note that if the bucket has more than 1000 objects,
    // we have to run this multiple times
    hasMoreObjects := true
    // Keep track of how many objects we delete
    totalObjects := 0

    for hasMoreObjects {
        resp, err := svc.ListObjects(&s3.ListObjectsInput{Bucket: aws.String(bucket)})
        if err != nil {
            exitErrorf("Unable to list items in bucket %q, %v", bucket, err)
        }

        numObjs := len(resp.Contents)
        totalObjects += numObjs

        // Create Delete object with slots for the objects to delete
        var items s3.Delete
        var objs= make([]*s3.ObjectIdentifier, numObjs)

        for i, o := range resp.Contents {
            // Add objects from command line to array
            objs[i] = &s3.ObjectIdentifier{Key: aws.String(*o.Key)}
        }

        // Add list of objects to delete to Delete object
        items.SetObjects(objs)

        // Delete the items
        _, err = svc.DeleteObjects(&s3.DeleteObjectsInput{Bucket: &bucket, Delete: &items})
        if err != nil {
            exitErrorf("Unable to delete objects from bucket %q, %v", bucket, err)
        }

        hasMoreObjects = *resp.IsTruncated
    }

    fmt.Println("Deleted", totalObjects, "object(s) from bucket", bucket)
}

func exitErrorf(msg string, args ...interface{}) {
    fmt.Fprintf(os.Stderr, msg+"\n", args...)
    os.Exit(1)
}
<commit_msg>Update example to use NewDeleteListIterator with Paginator for use with s3manager.NewBatchDelete.<commit_after>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 (
    "fmt"
    "os"

    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"
    "github.com/aws/aws-sdk-go/service/s3/s3manager"
)

// Deletes all of the objects in the specified S3 Bucket in the region configured in the shared config
// or AWS_REGION environment variable.
//
// Usage:
//    go run s3_delete_objects BUCKET
func main() {
    if len(os.Args) != 2 {
        exitErrorf("Bucket name required\nUsage: %s BUCKET", os.Args[0])
    }

    bucket := os.Args[1]

    // Initialize a session in us-west-2 that the SDK will use to load
    // credentials from the shared credentials file ~/.aws/credentials.
    sess, _ := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )

    // Create S3 service client
    svc := s3.New(sess)

    // Setup BatchDeleteIterator to iterate through a list of objects.
    iter := s3manager.NewDeleteListIterator(svc, &s3.ListObjectsInput{
        Bucket: aws.String(bucket),
    })

    // Traverse iterator deleting each object
    if err := s3manager.NewBatchDeleteWithClient(svc).Delete(aws.BackgroundContext(), iter); err != nil {
        exitErrorf("Unable to delete objects from bucket %q, %v", bucket, err)
    }

    fmt.Printf("Deleted object(s) from bucket: %s", bucket)
}

func exitErrorf(msg string, args ...interface{}) {
    fmt.Fprintf(os.Stderr, msg+"\n", args...)
    os.Exit(1)
}
<|endoftext|>


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

package cmd

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"sync"

	"github.com/cilium/cilium/pkg/controller"
	"github.com/cilium/cilium/pkg/datapath/connector"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/k8s"
	"github.com/cilium/cilium/pkg/k8s/watchers"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/maps/ctmap"
	"github.com/cilium/cilium/pkg/maps/lxcmap"
	"github.com/cilium/cilium/pkg/option"

	"github.com/sirupsen/logrus"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
)

type endpointRestoreState struct {
	restored []*endpoint.Endpoint
	toClean  []*endpoint.Endpoint
}

// validateEndpoint attempts to determine that the restored endpoint is valid, ie it
// still exists in k8s, its datapath devices are present, and Cilium is
// responsible for its workload, etc.
//
// Returns true to indicate that the endpoint is valid to restore, and an
// optional error.
func (d *Daemon) validateEndpoint(ep *endpoint.Endpoint) (valid bool, err error) {
	// On each restart, the health endpoint is supposed to be recreated.
	// Hence we need to clean health endpoint state unconditionally.
	if ep.HasLabels(labels.LabelHealth) {
		// Ignore health endpoint and don't report
		// it as not restored. But we need to clean up the old
		// state files, so do this now.
		healthStateDir := ep.StateDirectoryPath()
		scopedLog := log.WithFields(logrus.Fields{
			logfields.EndpointID: ep.ID,
			logfields.Path:       healthStateDir,
		})
		scopedLog.Debug("Removing old health endpoint state directory")
		if err := os.RemoveAll(healthStateDir); err != nil {
			scopedLog.Warning("Cannot clean up old health state directory")
		}
		return false, nil
	}

	if ep.K8sPodName != "" && ep.K8sNamespace != "" && k8s.IsEnabled() {
		d.k8sWatcher.WaitForCacheSync(watchers.K8sAPIGroupPodV1Core)
		_, err = d.k8sWatcher.GetCachedPod(ep.K8sNamespace, ep.K8sPodName)
		if err != nil && k8serrors.IsNotFound(err) {
			return false, fmt.Errorf("Kubernetes pod %s/%s does not exist", ep.K8sNamespace, ep.K8sPodName)
		}

		ep.RunMetadataResolver(d.fetchK8sLabelsAndAnnotations)
	}

	if err := ep.ValidateConnectorPlumbing(connector.CheckLink); err != nil {
		return false, err
	}

	if !ep.DatapathConfiguration.ExternalIpam {
		if err := d.allocateIPsLocked(ep); err != nil {
			return false, fmt.Errorf("Failed to re-allocate IP of endpoint: %s", err)
		}
	}

	return true, nil
}

// restoreOldEndpoints reads the list of existing endpoints previously managed by
// Cilium when it was last run and associated it with container workloads. This
// function performs the first step in restoring the endpoint structure,
// allocating their existing IPs out of the CIDR block and then inserting the
// endpoints into the endpoints list. It needs to be followed by a call to
// regenerateRestoredEndpoints() once the endpoint builder is ready.
//
// If clean is true, endpoints which cannot be associated with a container
// workloads are deleted.
func (d *Daemon) restoreOldEndpoints(dir string, clean bool) (*endpointRestoreState, error) {
	failed := 0
	state := &endpointRestoreState{
		restored: []*endpoint.Endpoint{},
		toClean:  []*endpoint.Endpoint{},
	}

	if !option.Config.RestoreState {
		log.Info("Endpoint restore is disabled, skipping restore step")
		return state, nil
	}

	log.Info("Restoring endpoints...")

	var (
		existingEndpoints map[string]*lxcmap.EndpointInfo
		err               error
	)

	if !option.Config.DryMode {
		existingEndpoints, err = lxcmap.DumpToMap()
		if err != nil {
			log.WithError(err).Warning("Unable to open endpoint map while restoring. Skipping cleanup of endpoint map on startup")
		}
	}

	dirFiles, err := ioutil.ReadDir(dir)
	if err != nil {
		return state, err
	}
	eptsID := endpoint.FilterEPDir(dirFiles)

	possibleEPs := endpoint.ReadEPsFromDirNames(d.ctx, d, dir, eptsID)

	if len(possibleEPs) == 0 {
		log.Info("No old endpoints found.")
		return state, nil
	}

	for _, ep := range possibleEPs {
		scopedLog := log.WithField(logfields.EndpointID, ep.ID)
		if k8s.IsEnabled() {
			scopedLog = scopedLog.WithField("k8sPodName", ep.GetK8sNamespaceAndPodName())
		}

		// We have to set the allocator for identities here during the Endpoint
		// lifecycle, because the identity allocator has been initialized *after*
		// endpoints are restored from disk. This is because we have to reserve
		// IPs for the endpoints that are restored via IPAM. Reserving of IPs
		// affects the allocation of IPs w.r.t. node addressing, which we need
		// to know before the identity allocator is initialized. We need to
		// know the node addressing because when adding a reference to the
		// kvstore because the local node's IP is used as a suffix for the key
		// in the key-value store.
		ep.SetAllocator(d.identityAllocator)

		restore, err := d.validateEndpoint(ep)
		if err != nil {
			// Disconnected EPs are not failures, clean them silently below
			if !ep.IsDisconnecting() {
				scopedLog.WithError(err).Warningf("Unable to restore endpoint, ignoring")
				failed++
			}
		}
		if !restore {
			if clean {
				state.toClean = append(state.toClean, ep)
			}
			continue
		}

		scopedLog.Debug("Restoring endpoint")
		ep.LogStatusOK(endpoint.Other, "Restoring endpoint from previous cilium instance")

		ep.SetDefaultConfiguration(true)
		ep.SetProxy(d.l7Proxy)
		ep.SkipStateClean()

		state.restored = append(state.restored, ep)

		if existingEndpoints != nil {
			delete(existingEndpoints, ep.IPv4.String())
			delete(existingEndpoints, ep.IPv6.String())
		}
	}

	log.WithFields(logrus.Fields{
		"restored": len(state.restored),
		"failed":   failed,
	}).Info("Endpoints restored")

	if existingEndpoints != nil {
		for hostIP, info := range existingEndpoints {
			if ip := net.ParseIP(hostIP); !info.IsHost() && ip != nil {
				if err := lxcmap.DeleteEntry(ip); err != nil {
					log.WithError(err).Warn("Unable to delete obsolete endpoint from BPF map")
				} else {
					log.Debugf("Removed outdated endpoint %d from endpoint map", info.LxcID)
				}
			}
		}
	}

	return state, nil
}

func (d *Daemon) regenerateRestoredEndpoints(state *endpointRestoreState) (restoreComplete chan struct{}) {
	restoreComplete = make(chan struct{}, 0)

	log.WithField("numRestored", len(state.restored)).Info("Regenerating restored endpoints")

	// Before regenerating, check whether the CT map has properties that
	// match this Cilium userspace instance. If not, it must be removed
	ctmap.DeleteIfUpgradeNeeded(nil)

	// we need to signalize when the endpoints are regenerated, i.e., when
	// they have finished to rebuild after being restored.
	epRegenerated := make(chan bool, len(state.restored))

	// Insert all endpoints into the endpoint list first before starting
	// the regeneration. This is required to ensure that if an individual
	// regeneration causes an identity change of an endpoint, the new
	// identity will trigger a policy recalculation of all endpoints to
	// account for the new identity during the grace period. For this
	// purpose, all endpoints being restored must already be in the
	// endpoint list.
	for i := len(state.restored) - 1; i >= 0; i-- {
		ep := state.restored[i]
		// If the endpoint has local conntrack option enabled, then
		// check whether the CT map needs upgrading (and do so).
		if ep.Options.IsEnabled(option.ConntrackLocal) {
			ctmap.DeleteIfUpgradeNeeded(ep)
		}

		// Insert into endpoint manager so it can be regenerated when calls to
		// RegenerateAllEndpoints() are made. This must be done synchronously (i.e.,
		// not in a goroutine) because regenerateRestoredEndpoints must guarantee
		// upon returning that endpoints are exposed to other subsystems via
		// endpointmanager.
		if err := ep.Expose(d.endpointManager); err != nil {
			log.WithError(err).Warning("Unable to restore endpoint")
			// remove endpoint from slice of endpoints to restore
			state.restored = append(state.restored[:i], state.restored[i+1:]...)
		}
	}

	for _, ep := range state.restored {
		log.WithField(logfields.EndpointID, ep.ID).Info("Successfully restored endpoint. Scheduling regeneration")
		go func(ep *endpoint.Endpoint, epRegenerated chan<- bool) {
			if err := ep.RegenerateAfterRestore(); err != nil {
				epRegenerated <- false
				return
			}
			epRegenerated <- true
		}(ep, epRegenerated)
	}

	var endpointCleanupCompleted sync.WaitGroup
	for _, ep := range state.toClean {
		endpointCleanupCompleted.Add(1)
		go func(ep *endpoint.Endpoint) {
			// The IP was not allocated yet so does not need to be free.
			// The identity may be allocated in the kvstore but we can't
			// release it easily as it will require to block on kvstore
			// connectivity which we can't do at this point. Let the lease
			// expire to release the identity.
			d.deleteEndpointQuiet(ep, endpoint.DeleteConfig{
				NoIdentityRelease: true,
				NoIPRelease:       true,
			})
			endpointCleanupCompleted.Done()
		}(ep)
	}
	endpointCleanupCompleted.Wait()

	go func() {
		regenerated, total := 0, 0
		if len(state.restored) > 0 {
			for buildSuccess := range epRegenerated {
				if buildSuccess {
					regenerated++
				}
				total++
				if total >= len(state.restored) {
					break
				}
			}
		}
		close(epRegenerated)

		log.WithFields(logrus.Fields{
			"regenerated": regenerated,
			"total":       total,
		}).Info("Finished regenerating restored endpoints")
		close(restoreComplete)
	}()

	return
}

func (d *Daemon) allocateIPsLocked(ep *endpoint.Endpoint) error {
	var err error

	if option.Config.EnableIPv6 && ep.IPv6 != nil {
		err = d.ipam.AllocateIPWithoutSyncUpstream(ep.IPv6.IP(), ep.HumanStringLocked()+" [restored]")
		if err != nil {
			return fmt.Errorf("unable to reallocate %s IPv6 address: %s", ep.IPv6.IP(), err)
		}

		defer func() {
			if err != nil {
				d.ipam.ReleaseIP(ep.IPv6.IP())
			}
		}()
	}

	if option.Config.EnableIPv4 && ep.IPv4 != nil {
		if err = d.ipam.AllocateIPWithoutSyncUpstream(ep.IPv4.IP(), ep.HumanStringLocked()+" [restored]"); err != nil {
			return fmt.Errorf("unable to reallocate %s IPv4 address: %s", ep.IPv4.IP(), err)
		}
	}

	return nil
}

func (d *Daemon) initRestore(restoredEndpoints *endpointRestoreState) chan struct{} {
	bootstrapStats.restore.Start()
	var restoreComplete chan struct{}
	if option.Config.RestoreState {
		// When we regenerate restored endpoints, it is guaranteed tha we have
		// received the full list of policies present at the time the daemon
		// is bootstrapped.
		restoreComplete = d.regenerateRestoredEndpoints(restoredEndpoints)
		go func() {
			<-restoreComplete
			endParallelMapMode()
		}()

		go func() {
			if k8s.IsEnabled() {
				// Also wait for all cluster mesh to be synchronized with the
				// datapath before proceeding.
				if option.Config.ClusterID != 0 {
					err := d.clustermesh.ClustersSynced(context.Background())
					if err != nil {
						log.WithError(err).Fatal("timeout while waiting for all clusters to be locally synchronized")
					}
				}
				// Start controller which removes any leftover Kubernetes
				// services that may have been deleted while Cilium was not
				// running. Once this controller succeeds, because it has no
				// RunInterval specified, it will not run again unless updated
				// elsewhere. This means that if, for instance, a user manually
				// adds a service via the CLI into the BPF maps, that it will
				// not be cleaned up by the daemon until it restarts.
				controller.NewManager().UpdateController("sync-lb-maps-with-k8s-services",
					controller.ControllerParams{
						DoFunc: func(ctx context.Context) error {
							return d.svc.SyncWithK8sFinished()
						},
						Context: d.ctx,
					},
				)
			}
		}()
	} else {
		log.Info("State restore is disabled. Existing endpoints on node are ignored")

		// No restore happened, end parallel map mode immediately
		endParallelMapMode()
	}
	bootstrapStats.restore.End(true)

	return restoreComplete
}
<commit_msg>agent: fix panic when clustermesh not set and cluster-id is non-zero<commit_after>// Copyright 2016-2020 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"sync"

	"github.com/cilium/cilium/pkg/controller"
	"github.com/cilium/cilium/pkg/datapath/connector"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/k8s"
	"github.com/cilium/cilium/pkg/k8s/watchers"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/maps/ctmap"
	"github.com/cilium/cilium/pkg/maps/lxcmap"
	"github.com/cilium/cilium/pkg/option"

	"github.com/sirupsen/logrus"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
)

type endpointRestoreState struct {
	restored []*endpoint.Endpoint
	toClean  []*endpoint.Endpoint
}

// validateEndpoint attempts to determine that the restored endpoint is valid, ie it
// still exists in k8s, its datapath devices are present, and Cilium is
// responsible for its workload, etc.
//
// Returns true to indicate that the endpoint is valid to restore, and an
// optional error.
func (d *Daemon) validateEndpoint(ep *endpoint.Endpoint) (valid bool, err error) {
	// On each restart, the health endpoint is supposed to be recreated.
	// Hence we need to clean health endpoint state unconditionally.
	if ep.HasLabels(labels.LabelHealth) {
		// Ignore health endpoint and don't report
		// it as not restored. But we need to clean up the old
		// state files, so do this now.
		healthStateDir := ep.StateDirectoryPath()
		scopedLog := log.WithFields(logrus.Fields{
			logfields.EndpointID: ep.ID,
			logfields.Path:       healthStateDir,
		})
		scopedLog.Debug("Removing old health endpoint state directory")
		if err := os.RemoveAll(healthStateDir); err != nil {
			scopedLog.Warning("Cannot clean up old health state directory")
		}
		return false, nil
	}

	if ep.K8sPodName != "" && ep.K8sNamespace != "" && k8s.IsEnabled() {
		d.k8sWatcher.WaitForCacheSync(watchers.K8sAPIGroupPodV1Core)
		_, err = d.k8sWatcher.GetCachedPod(ep.K8sNamespace, ep.K8sPodName)
		if err != nil && k8serrors.IsNotFound(err) {
			return false, fmt.Errorf("Kubernetes pod %s/%s does not exist", ep.K8sNamespace, ep.K8sPodName)
		}

		ep.RunMetadataResolver(d.fetchK8sLabelsAndAnnotations)
	}

	if err := ep.ValidateConnectorPlumbing(connector.CheckLink); err != nil {
		return false, err
	}

	if !ep.DatapathConfiguration.ExternalIpam {
		if err := d.allocateIPsLocked(ep); err != nil {
			return false, fmt.Errorf("Failed to re-allocate IP of endpoint: %s", err)
		}
	}

	return true, nil
}

// restoreOldEndpoints reads the list of existing endpoints previously managed by
// Cilium when it was last run and associated it with container workloads. This
// function performs the first step in restoring the endpoint structure,
// allocating their existing IPs out of the CIDR block and then inserting the
// endpoints into the endpoints list. It needs to be followed by a call to
// regenerateRestoredEndpoints() once the endpoint builder is ready.
//
// If clean is true, endpoints which cannot be associated with a container
// workloads are deleted.
func (d *Daemon) restoreOldEndpoints(dir string, clean bool) (*endpointRestoreState, error) {
	failed := 0
	state := &endpointRestoreState{
		restored: []*endpoint.Endpoint{},
		toClean:  []*endpoint.Endpoint{},
	}

	if !option.Config.RestoreState {
		log.Info("Endpoint restore is disabled, skipping restore step")
		return state, nil
	}

	log.Info("Restoring endpoints...")

	var (
		existingEndpoints map[string]*lxcmap.EndpointInfo
		err               error
	)

	if !option.Config.DryMode {
		existingEndpoints, err = lxcmap.DumpToMap()
		if err != nil {
			log.WithError(err).Warning("Unable to open endpoint map while restoring. Skipping cleanup of endpoint map on startup")
		}
	}

	dirFiles, err := ioutil.ReadDir(dir)
	if err != nil {
		return state, err
	}
	eptsID := endpoint.FilterEPDir(dirFiles)

	possibleEPs := endpoint.ReadEPsFromDirNames(d.ctx, d, dir, eptsID)

	if len(possibleEPs) == 0 {
		log.Info("No old endpoints found.")
		return state, nil
	}

	for _, ep := range possibleEPs {
		scopedLog := log.WithField(logfields.EndpointID, ep.ID)
		if k8s.IsEnabled() {
			scopedLog = scopedLog.WithField("k8sPodName", ep.GetK8sNamespaceAndPodName())
		}

		// We have to set the allocator for identities here during the Endpoint
		// lifecycle, because the identity allocator has been initialized *after*
		// endpoints are restored from disk. This is because we have to reserve
		// IPs for the endpoints that are restored via IPAM. Reserving of IPs
		// affects the allocation of IPs w.r.t. node addressing, which we need
		// to know before the identity allocator is initialized. We need to
		// know the node addressing because when adding a reference to the
		// kvstore because the local node's IP is used as a suffix for the key
		// in the key-value store.
		ep.SetAllocator(d.identityAllocator)

		restore, err := d.validateEndpoint(ep)
		if err != nil {
			// Disconnected EPs are not failures, clean them silently below
			if !ep.IsDisconnecting() {
				scopedLog.WithError(err).Warningf("Unable to restore endpoint, ignoring")
				failed++
			}
		}
		if !restore {
			if clean {
				state.toClean = append(state.toClean, ep)
			}
			continue
		}

		scopedLog.Debug("Restoring endpoint")
		ep.LogStatusOK(endpoint.Other, "Restoring endpoint from previous cilium instance")

		ep.SetDefaultConfiguration(true)
		ep.SetProxy(d.l7Proxy)
		ep.SkipStateClean()

		state.restored = append(state.restored, ep)

		if existingEndpoints != nil {
			delete(existingEndpoints, ep.IPv4.String())
			delete(existingEndpoints, ep.IPv6.String())
		}
	}

	log.WithFields(logrus.Fields{
		"restored": len(state.restored),
		"failed":   failed,
	}).Info("Endpoints restored")

	if existingEndpoints != nil {
		for hostIP, info := range existingEndpoints {
			if ip := net.ParseIP(hostIP); !info.IsHost() && ip != nil {
				if err := lxcmap.DeleteEntry(ip); err != nil {
					log.WithError(err).Warn("Unable to delete obsolete endpoint from BPF map")
				} else {
					log.Debugf("Removed outdated endpoint %d from endpoint map", info.LxcID)
				}
			}
		}
	}

	return state, nil
}

func (d *Daemon) regenerateRestoredEndpoints(state *endpointRestoreState) (restoreComplete chan struct{}) {
	restoreComplete = make(chan struct{}, 0)

	log.WithField("numRestored", len(state.restored)).Info("Regenerating restored endpoints")

	// Before regenerating, check whether the CT map has properties that
	// match this Cilium userspace instance. If not, it must be removed
	ctmap.DeleteIfUpgradeNeeded(nil)

	// we need to signalize when the endpoints are regenerated, i.e., when
	// they have finished to rebuild after being restored.
	epRegenerated := make(chan bool, len(state.restored))

	// Insert all endpoints into the endpoint list first before starting
	// the regeneration. This is required to ensure that if an individual
	// regeneration causes an identity change of an endpoint, the new
	// identity will trigger a policy recalculation of all endpoints to
	// account for the new identity during the grace period. For this
	// purpose, all endpoints being restored must already be in the
	// endpoint list.
	for i := len(state.restored) - 1; i >= 0; i-- {
		ep := state.restored[i]
		// If the endpoint has local conntrack option enabled, then
		// check whether the CT map needs upgrading (and do so).
		if ep.Options.IsEnabled(option.ConntrackLocal) {
			ctmap.DeleteIfUpgradeNeeded(ep)
		}

		// Insert into endpoint manager so it can be regenerated when calls to
		// RegenerateAllEndpoints() are made. This must be done synchronously (i.e.,
		// not in a goroutine) because regenerateRestoredEndpoints must guarantee
		// upon returning that endpoints are exposed to other subsystems via
		// endpointmanager.
		if err := ep.Expose(d.endpointManager); err != nil {
			log.WithError(err).Warning("Unable to restore endpoint")
			// remove endpoint from slice of endpoints to restore
			state.restored = append(state.restored[:i], state.restored[i+1:]...)
		}
	}

	for _, ep := range state.restored {
		log.WithField(logfields.EndpointID, ep.ID).Info("Successfully restored endpoint. Scheduling regeneration")
		go func(ep *endpoint.Endpoint, epRegenerated chan<- bool) {
			if err := ep.RegenerateAfterRestore(); err != nil {
				epRegenerated <- false
				return
			}
			epRegenerated <- true
		}(ep, epRegenerated)
	}

	var endpointCleanupCompleted sync.WaitGroup
	for _, ep := range state.toClean {
		endpointCleanupCompleted.Add(1)
		go func(ep *endpoint.Endpoint) {
			// The IP was not allocated yet so does not need to be free.
			// The identity may be allocated in the kvstore but we can't
			// release it easily as it will require to block on kvstore
			// connectivity which we can't do at this point. Let the lease
			// expire to release the identity.
			d.deleteEndpointQuiet(ep, endpoint.DeleteConfig{
				NoIdentityRelease: true,
				NoIPRelease:       true,
			})
			endpointCleanupCompleted.Done()
		}(ep)
	}
	endpointCleanupCompleted.Wait()

	go func() {
		regenerated, total := 0, 0
		if len(state.restored) > 0 {
			for buildSuccess := range epRegenerated {
				if buildSuccess {
					regenerated++
				}
				total++
				if total >= len(state.restored) {
					break
				}
			}
		}
		close(epRegenerated)

		log.WithFields(logrus.Fields{
			"regenerated": regenerated,
			"total":       total,
		}).Info("Finished regenerating restored endpoints")
		close(restoreComplete)
	}()

	return
}

func (d *Daemon) allocateIPsLocked(ep *endpoint.Endpoint) error {
	var err error

	if option.Config.EnableIPv6 && ep.IPv6 != nil {
		err = d.ipam.AllocateIPWithoutSyncUpstream(ep.IPv6.IP(), ep.HumanStringLocked()+" [restored]")
		if err != nil {
			return fmt.Errorf("unable to reallocate %s IPv6 address: %s", ep.IPv6.IP(), err)
		}

		defer func() {
			if err != nil {
				d.ipam.ReleaseIP(ep.IPv6.IP())
			}
		}()
	}

	if option.Config.EnableIPv4 && ep.IPv4 != nil {
		if err = d.ipam.AllocateIPWithoutSyncUpstream(ep.IPv4.IP(), ep.HumanStringLocked()+" [restored]"); err != nil {
			return fmt.Errorf("unable to reallocate %s IPv4 address: %s", ep.IPv4.IP(), err)
		}
	}

	return nil
}

func (d *Daemon) initRestore(restoredEndpoints *endpointRestoreState) chan struct{} {
	bootstrapStats.restore.Start()
	var restoreComplete chan struct{}
	if option.Config.RestoreState {
		// When we regenerate restored endpoints, it is guaranteed tha we have
		// received the full list of policies present at the time the daemon
		// is bootstrapped.
		restoreComplete = d.regenerateRestoredEndpoints(restoredEndpoints)
		go func() {
			<-restoreComplete
			endParallelMapMode()
		}()

		go func() {
			if k8s.IsEnabled() {
				// Also wait for all cluster mesh to be synchronized with the
				// datapath before proceeding.
				if d.clustermesh != nil {
					err := d.clustermesh.ClustersSynced(context.Background())
					if err != nil {
						log.WithError(err).Fatal("timeout while waiting for all clusters to be locally synchronized")
					}
				}
				// Start controller which removes any leftover Kubernetes
				// services that may have been deleted while Cilium was not
				// running. Once this controller succeeds, because it has no
				// RunInterval specified, it will not run again unless updated
				// elsewhere. This means that if, for instance, a user manually
				// adds a service via the CLI into the BPF maps, that it will
				// not be cleaned up by the daemon until it restarts.
				controller.NewManager().UpdateController("sync-lb-maps-with-k8s-services",
					controller.ControllerParams{
						DoFunc: func(ctx context.Context) error {
							return d.svc.SyncWithK8sFinished()
						},
						Context: d.ctx,
					},
				)
			}
		}()
	} else {
		log.Info("State restore is disabled. Existing endpoints on node are ignored")

		// No restore happened, end parallel map mode immediately
		endParallelMapMode()
	}
	bootstrapStats.restore.End(true)

	return restoreComplete
}
<|endoftext|>


<commit_before>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main
/*
 *  Filename:    godirs.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Tue Jul  5 22:13:49 PDT 2011
 *  Description:
 *  Usage:       godirs [options] ARGUMENT ...
 */
import (
    "os"
    "fmt"
    "flag"
    "sync"
    "time"
    "log"
    //"strings"
    "github.com/bmatsuo/dispatch"
    "path/filepath"
)
const (
    stdDelay = 50e3
)

type Options struct {
    list     func(string) ([]string,[]int64)
    beClever bool
    dir      string
    verbose  bool
}
var opt = Options{}
func SetupFlags() *flag.FlagSet {
    var fs = flag.NewFlagSet("godirs", flag.ExitOnError)
    fs.BoolVar(&(opt.beClever), "c", false, "Use the clever Dispatch.")
    fs.BoolVar(&(opt.verbose), "v", false, "Verbose program output.")
    return fs
}
func VerifyFlags(fs *flag.FlagSet) {
    if fs.NArg() < 1 {
        fmt.Fprintf(os.Stderr, "missing DIR argument.")
    }
}
func ParseFlags() {
    var fs = SetupFlags()
    fs.Parse(os.Args[1:])
    VerifyFlags(fs)
    opt.dir = fs.Arg(0)
    if opt.beClever {
        opt.list = WalkerList
    } else {
        opt.list = StupidWalkerList
    }
}

type Walker struct {
    gq     *dispatch.Dispatch
    done   chan bool
    paths  []string
    sizes  []int64
    lock   *sync.Mutex
    wg     *sync.WaitGroup
}
func NewWalker() *Walker {
    var w    = new(Walker)
    w.gq     = dispatch.New(20)
    w.lock   = new(sync.Mutex)
    w.wg     = new(sync.WaitGroup)
    w.done   = make(chan bool)
    w.paths  = make([]string, 0, 1)
    w.sizes  = make([]int64, 0, 1)
    return w
}
func (w *Walker) VisitFile(path string, info *os.FileInfo) {
    var f = func (id int64) {
        log.Print("stating %s", path)
        var stat, err = os.Stat(path)
        if err != nil {
            panic(err)
        }
        time.Sleep(stdDelay)
        w.lock.Lock()
        w.sizes = append(w.sizes, stat.Size)
        w.paths = append(w.paths, path)
        w.lock.Unlock()
        w.wg.Done()
    }
    w.wg.Add(1)
    w.gq.Enqueue(&dispatch.StdTask{f})
}
func (w *Walker) VisitDir(path string, f *os.FileInfo) bool {
    return true
}
func WalkerList(dir string) ([]string, []int64) {
    var w = NewWalker()
    var errors = make(chan os.Error)
    go func() {
        for e := range errors {
            panic("Walk error: " + e.String())
        }
    } ()
    go w.gq.Start()
    filepath.Walk(dir, w, errors)
    w.wg.Wait()
    return w.paths, w.sizes
}

type StupidWalker struct {
    paths  []string
    sizes  []int64
}
func NewStupidWalker() *StupidWalker {
    var sw    = new(StupidWalker)
    sw.paths  = make([]string, 0, 1)
    sw.sizes  = make([]int64, 0, 1)
    return sw
}
func (sw *StupidWalker) VisitFile(path string, info *os.FileInfo) {
    var stat, err = os.Stat(path)
    if err != nil {
        panic(err)
    }
    time.Sleep(stdDelay)
    sw.sizes = append(sw.sizes, stat.Size)
    sw.paths = append(sw.paths, path)
}
func (sw *StupidWalker) VisitDir(path string, f *os.FileInfo) bool {
    return true
}
func StupidWalkerList(dir string) ([]string, []int64) {
    var sw = NewStupidWalker()
    var errors = make(chan os.Error)
    var done = make(chan bool)
    go func() {
        for e := range errors {
            panic("Walk error: " + e.String())
        }
        done<-true
    } ()
    filepath.Walk(dir, sw, errors)
    close(errors)
    <-done
    return sw.paths, sw.sizes
}

func main() {
    ParseFlags()
    var paths, sizes = opt.list(opt.dir)
    for i, path := range paths {
        fmt.Printf("\t%50s %d\n", path, sizes[i])
    }
}

/*
type recLister struct {
    BaseWait   int64
    MaxProc    int
    waittime   int64
    maxwait    int64
    processing int
    mutex      *sync.Mutex
    errors     chan os.Error
    done       chan bool
    paths      []string
    sizes      []int64
}
func newRecLister() *recLister {
    var rl = new(recLister)
    rl.mutex    = new(sync.Mutex)
    rl.errors   = make(chan os.Error)
    rl.done     = make(chan bool)
    rl.paths    = make([]string, 0, 1)
    rl.sizes    = make([]int64, 0, 1)
    rl.MaxProc  = 5
    rl.BaseWait = 10
    rl.waittime = rl.BaseWait
    rl.maxwait  = 500e6
    return rl
}
func (rl *recLister) VisitDir(path string, f *os.FileInfo) bool {
    //rl.paths = append(rl.paths, path)
    return true
}
func (rl *recLister) visitFile(path string, f *os.FileInfo) {
    var stat, err = os.Stat(path)
    if err != nil {
        panic(err)
    }
    rl.sizes = append(rl.sizes, stat.Size)
    rl.paths = append(rl.paths, path)
    //time.Sleep(5e9)
    rl.mutex.Lock()
    rl.processing--
    rl.mutex.Unlock()
}
func (rl *recLister) VisitFile(path string, f *os.FileInfo) {
    for true {
        // Attempt to start processing the file.
        rl.mutex.Lock()
        if rl.processing >= rl.MaxProc {
            // Too many threads, wait and try again.
            rl.waittime <<= 2
            if rl.waittime > rl.maxwait {
                rl.waittime = rl.maxwait
            }
            rl.mutex.Unlock()
            time.Sleep(rl.waittime)
            continue
        }
        // Keep the books and reset wait time before unlocking and processing.
        rl.processing++
        rl.waittime = rl.BaseWait
        rl.mutex.Unlock()
        go rl.visitFile(path, f)
        return
    }
}

func RecFileList(dir string) ([]string, []int64) {
    var rl = newRecLister()
    go func() {
        for e := range rl.errors {
            panic("Walk error: " + e.String())
        }
        rl.done <- true
    } ()
    filepath.Walk(dir, rl, rl.errors)
    close(rl.errors)
    <-rl.done
    return rl.paths, rl.sizes
}
*/
<commit_msg>Stop logging output in godu<commit_after>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main
/*
 *  Filename:    godirs.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Tue Jul  5 22:13:49 PDT 2011
 *  Description:
 *  Usage:       godirs [options] ARGUMENT ...
 */
import (
    "os"
    "fmt"
    "flag"
    "sync"
    "time"
    //"log"
    //"strings"
    "github.com/bmatsuo/dispatch"
    "path/filepath"
)
const (
    stdDelay = 50e3
)

type Options struct {
    list     func(string) ([]string,[]int64)
    beClever bool
    dir      string
    verbose  bool
}
var opt = Options{}
func SetupFlags() *flag.FlagSet {
    var fs = flag.NewFlagSet("godirs", flag.ExitOnError)
    fs.BoolVar(&(opt.beClever), "c", false, "Use the clever Dispatch.")
    fs.BoolVar(&(opt.verbose), "v", false, "Verbose program output.")
    return fs
}
func VerifyFlags(fs *flag.FlagSet) {
    if fs.NArg() < 1 {
        fmt.Fprintf(os.Stderr, "missing DIR argument.")
    }
}
func ParseFlags() {
    var fs = SetupFlags()
    fs.Parse(os.Args[1:])
    VerifyFlags(fs)
    opt.dir = fs.Arg(0)
    if opt.beClever {
        opt.list = WalkerList
    } else {
        opt.list = StupidWalkerList
    }
}

type Walker struct {
    gq     *dispatch.Dispatch
    done   chan bool
    paths  []string
    sizes  []int64
    lock   *sync.Mutex
    wg     *sync.WaitGroup
}
func NewWalker() *Walker {
    var w    = new(Walker)
    w.gq     = dispatch.New(20)
    w.lock   = new(sync.Mutex)
    w.wg     = new(sync.WaitGroup)
    w.done   = make(chan bool)
    w.paths  = make([]string, 0, 1)
    w.sizes  = make([]int64, 0, 1)
    return w
}
func (w *Walker) VisitFile(path string, info *os.FileInfo) {
    var f = func (id int64) {
        //log.Print("stating %s", path)
        var stat, err = os.Stat(path)
        if err != nil {
            panic(err)
        }
        time.Sleep(stdDelay)
        w.lock.Lock()
        w.sizes = append(w.sizes, stat.Size)
        w.paths = append(w.paths, path)
        w.lock.Unlock()
        w.wg.Done()
    }
    w.wg.Add(1)
    w.gq.Enqueue(&dispatch.StdTask{f})
}
func (w *Walker) VisitDir(path string, f *os.FileInfo) bool {
    return true
}
func WalkerList(dir string) ([]string, []int64) {
    var w = NewWalker()
    var errors = make(chan os.Error)
    go func() {
        for e := range errors {
            panic("Walk error: " + e.String())
        }
    } ()
    go w.gq.Start()
    filepath.Walk(dir, w, errors)
    w.wg.Wait()
    return w.paths, w.sizes
}

type StupidWalker struct {
    paths  []string
    sizes  []int64
}
func NewStupidWalker() *StupidWalker {
    var sw    = new(StupidWalker)
    sw.paths  = make([]string, 0, 1)
    sw.sizes  = make([]int64, 0, 1)
    return sw
}
func (sw *StupidWalker) VisitFile(path string, info *os.FileInfo) {
    var stat, err = os.Stat(path)
    if err != nil {
        panic(err)
    }
    time.Sleep(stdDelay)
    sw.sizes = append(sw.sizes, stat.Size)
    sw.paths = append(sw.paths, path)
}
func (sw *StupidWalker) VisitDir(path string, f *os.FileInfo) bool {
    return true
}
func StupidWalkerList(dir string) ([]string, []int64) {
    var sw = NewStupidWalker()
    var errors = make(chan os.Error)
    var done = make(chan bool)
    go func() {
        for e := range errors {
            panic("Walk error: " + e.String())
        }
        done<-true
    } ()
    filepath.Walk(dir, sw, errors)
    close(errors)
    <-done
    return sw.paths, sw.sizes
}

func main() {
    ParseFlags()
    var paths, sizes = opt.list(opt.dir)
    for i, path := range paths {
        fmt.Printf("\t%50s %d\n", path, sizes[i])
    }
}

/*
type recLister struct {
    BaseWait   int64
    MaxProc    int
    waittime   int64
    maxwait    int64
    processing int
    mutex      *sync.Mutex
    errors     chan os.Error
    done       chan bool
    paths      []string
    sizes      []int64
}
func newRecLister() *recLister {
    var rl = new(recLister)
    rl.mutex    = new(sync.Mutex)
    rl.errors   = make(chan os.Error)
    rl.done     = make(chan bool)
    rl.paths    = make([]string, 0, 1)
    rl.sizes    = make([]int64, 0, 1)
    rl.MaxProc  = 5
    rl.BaseWait = 10
    rl.waittime = rl.BaseWait
    rl.maxwait  = 500e6
    return rl
}
func (rl *recLister) VisitDir(path string, f *os.FileInfo) bool {
    //rl.paths = append(rl.paths, path)
    return true
}
func (rl *recLister) visitFile(path string, f *os.FileInfo) {
    var stat, err = os.Stat(path)
    if err != nil {
        panic(err)
    }
    rl.sizes = append(rl.sizes, stat.Size)
    rl.paths = append(rl.paths, path)
    //time.Sleep(5e9)
    rl.mutex.Lock()
    rl.processing--
    rl.mutex.Unlock()
}
func (rl *recLister) VisitFile(path string, f *os.FileInfo) {
    for true {
        // Attempt to start processing the file.
        rl.mutex.Lock()
        if rl.processing >= rl.MaxProc {
            // Too many threads, wait and try again.
            rl.waittime <<= 2
            if rl.waittime > rl.maxwait {
                rl.waittime = rl.maxwait
            }
            rl.mutex.Unlock()
            time.Sleep(rl.waittime)
            continue
        }
        // Keep the books and reset wait time before unlocking and processing.
        rl.processing++
        rl.waittime = rl.BaseWait
        rl.mutex.Unlock()
        go rl.visitFile(path, f)
        return
    }
}

func RecFileList(dir string) ([]string, []int64) {
    var rl = newRecLister()
    go func() {
        for e := range rl.errors {
            panic("Walk error: " + e.String())
        }
        rl.done <- true
    } ()
    filepath.Walk(dir, rl, rl.errors)
    close(rl.errors)
    <-rl.done
    return rl.paths, rl.sizes
}
*/
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io/ioutil"
	"testing"
)

func setup() {
}

func teardown() {
}

func TestAppWithoutApiKey(t *testing.T) {
	app := buildApp()
	app.Writer = ioutil.Discard

	// test with other global flags
	tests := []struct {
		args    []string
		wantErr error
	}{
		{
			args:    []string{"doctl"},
			wantErr: errors.New("must provide API Key via DIGITALOCEAN_API_KEY environment variable or via CLI argument."),
		},
		{
			args:    []string{"doctl", "--version"},
			wantErr: nil,
		},
		{
			args:    []string{"doctl", "--help"},
			wantErr: nil,
		},
	}

	for _, tt := range tests {
		err := app.Run(tt.args)
		if tt.wantErr == nil {
			if err != nil {
				t.Errorf("Expected nil error")
			}
		} else if err.Error() != tt.wantErr.Error() {
			t.Errorf("app.Run(%v) = %#v, want %#v", tt.args, err, tt.wantErr)
		}
	}
}

func TestGlobalFormatFlag(t *testing.T) {
	app := buildApp()
	app.Writer = ioutil.Discard

	args := []string{"doctl", "-k", "key", "-f", "invalid"}
	err := app.Run(args)

	expected := `invalid output format: "invalid", available output options: json, yaml.`
	if err.Error() != expected {
		t.Errorf("app.Run(%v) = %#v, want %#v", err, err.Error(), expected)
	}
}
<commit_msg>Update test to reflect non-error<commit_after>package main

import (
	"io/ioutil"
	"testing"
)

func setup() {
}

func teardown() {
}

func TestAppWithoutApiKey(t *testing.T) {
	app := buildApp()
	app.Writer = ioutil.Discard

	// test with other global flags
	tests := []struct {
		args    []string
		wantErr error
	}{
		{
			args:    []string{"doctl"},
			wantErr: nil,
		},
		{
			args:    []string{"doctl", "--version"},
			wantErr: nil,
		},
		{
			args:    []string{"doctl", "--help"},
			wantErr: nil,
		},
	}

	for _, tt := range tests {
		err := app.Run(tt.args)
		if tt.wantErr == nil {
			if err != nil {
				t.Errorf("Expected nil error")
			}
		} else if err.Error() != tt.wantErr.Error() {
			t.Errorf("app.Run(%v) = %#v, want %#v", tt.args, err, tt.wantErr)
		}
	}
}

func TestGlobalFormatFlag(t *testing.T) {
	app := buildApp()
	app.Writer = ioutil.Discard

	args := []string{"doctl", "-k", "key", "-f", "invalid"}
	err := app.Run(args)

	expected := `invalid output format: "invalid", available output options: json, yaml.`
	if err.Error() != expected {
		t.Errorf("app.Run(%v) = %#v, want %#v", err, err.Error(), expected)
	}
}
<|endoftext|>


<commit_before>package gohex

import (
	"testing"
)

func TestConstructor(t *testing.T) {
	m := NewMemory()
	if m.GetStartAddress() != 0 {
		t.Errorf("wrong initial start address")
	}
	if len(m.GetDataSegments()) != 0 {
		t.Errorf("wrong initial data segments")
	}
}

func assertParseError(t *testing.T, m *Memory, input string, et ParseErrorType, err string) {
	if e := m.ParseIntelHex(input); e != nil {
		perr, ok := e.(*ParseError)
		if ok == true {
			if perr.ErrorType != et {
				t.Errorf(perr.Error())
				t.Errorf(err)
			}
		} else {
			t.Errorf(err)
		}
	} else {
		t.Errorf(err)
	}
}

func TestSyntaxError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, "00000001FF\n", SYNTAX_ERROR, "no colon error")
	assertParseError(t, m, ":qw00000001FF\n", SYNTAX_ERROR, "no ascii hex error")
	assertParseError(t, m, ":0000001FF\n", SYNTAX_ERROR, "no odd/even hex error")
}

func TestDataError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, ":000000FF\n", DATA_ERROR, "no line length error")
	assertParseError(t, m, ":02000000FE\n", DATA_ERROR, "no data length error")
	assertParseError(t, m, "\n", DATA_ERROR, "no end of file line error")
	assertParseError(t, m, ":000000FF01\n", DATA_ERROR, "no end of file line error")
	assertParseError(t, m, ":0400000501000000F6\n", DATA_ERROR, "no end of file line error")
	assertParseError(t, m, ":0400000501000000F6\n:0400000502000000F5\n:00000001FF\n", DATA_ERROR, "no multiple start address lines error")
}

func TestChecksumError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, ":00000101FF\n", CHECKSUM_ERROR, "no checking checksum error")
	assertParseError(t, m, ":00000001FE\n", CHECKSUM_ERROR, "no checking checksum error")
	assertParseError(t, m, ":0000000001\n", CHECKSUM_ERROR, "no checking checksum error")
	assertParseError(t, m, ":000000FF02\n", CHECKSUM_ERROR, "no checking checksum error")
}

func TestRecordsError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, ":00000101FE\n", RECORD_ERROR, "no eof record error")
	assertParseError(t, m, ":00010001FE\n", RECORD_ERROR, "no eof record error")
	assertParseError(t, m, ":0100000100FE\n", RECORD_ERROR, "no eof record error")
	assertParseError(t, m, ":020001040101F7\n", RECORD_ERROR, "no extended address record error")
	assertParseError(t, m, ":020100040101F7\n", RECORD_ERROR, "no extended address record error")
	assertParseError(t, m, ":03000004010100F7\n", RECORD_ERROR, "no extended address record error")
	assertParseError(t, m, ":0400010501010101F2\n", RECORD_ERROR, "no start address record error")
	assertParseError(t, m, ":0401000501010101F2\n", RECORD_ERROR, "no start address record error")
	assertParseError(t, m, ":050000050101010100F2\n", RECORD_ERROR, "no start address record error")
}
<commit_msg>more tests<commit_after>package gohex

import (
	"testing"
)

func TestConstructor(t *testing.T) {
	m := NewMemory()
	if m.GetStartAddress() != 0 {
		t.Error("incorrect initial start address")
	}
	if len(m.GetDataSegments()) != 0 {
		t.Error("incorrect initial data segments")
	}
	if m.currentAddress != 0 {
		t.Error("incorrect initial data segments")
	}
}

func assertParseError(t *testing.T, m *Memory, input string, et ParseErrorType, err string) {
	if e := m.ParseIntelHex(input); e != nil {
		perr, ok := e.(*ParseError)
		if ok == true {
			if perr.ErrorType != et {
				t.Error(perr.Error())
				t.Error(err)
			}
		} else {
			t.Error(err)
		}
	} else {
		t.Error(err)
	}
}

func TestSyntaxError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, "00000001FF\n", SYNTAX_ERROR, "no colon error")
	assertParseError(t, m, ":qw00000001FF\n", SYNTAX_ERROR, "no ascii hex error")
	assertParseError(t, m, ":0000001FF\n", SYNTAX_ERROR, "no odd/even hex error")
}

func TestDataError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, ":000000FF\n", DATA_ERROR, "no line length error")
	assertParseError(t, m, ":02000000FE\n", DATA_ERROR, "no data length error")
	assertParseError(t, m, "\n", DATA_ERROR, "no end of file line error")
	assertParseError(t, m, ":000000FF01\n", DATA_ERROR, "no end of file line error")
	assertParseError(t, m, ":0400000501000000F6\n", DATA_ERROR, "no end of file line error")
	assertParseError(t, m, ":0400000501000000F6\n:0400000502000000F5\n:00000001FF\n", DATA_ERROR, "no multiple start address lines error")
}

func TestChecksumError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, ":00000101FF\n", CHECKSUM_ERROR, "no checking checksum error")
	assertParseError(t, m, ":00000001FE\n", CHECKSUM_ERROR, "no checking checksum error")
	assertParseError(t, m, ":0000000001\n", CHECKSUM_ERROR, "no checking checksum error")
	assertParseError(t, m, ":000000FF02\n", CHECKSUM_ERROR, "no checking checksum error")
}

func TestRecordsError(t *testing.T) {
	m := NewMemory()
	assertParseError(t, m, ":00000101FE\n", RECORD_ERROR, "no eof record error")
	assertParseError(t, m, ":00010001FE\n", RECORD_ERROR, "no eof record error")
	assertParseError(t, m, ":0100000100FE\n", RECORD_ERROR, "no eof record error")
	assertParseError(t, m, ":020001040101F7\n", RECORD_ERROR, "no extended address record error")
	assertParseError(t, m, ":020100040101F7\n", RECORD_ERROR, "no extended address record error")
	assertParseError(t, m, ":03000004010100F7\n", RECORD_ERROR, "no extended address record error")
	assertParseError(t, m, ":0400010501010101F2\n", RECORD_ERROR, "no start address record error")
	assertParseError(t, m, ":0401000501010101F2\n", RECORD_ERROR, "no start address record error")
	assertParseError(t, m, ":050000050101010100F2\n", RECORD_ERROR, "no start address record error")
}

func TestAddressError(t *testing.T) {
	m := NewMemory()
	err := m.ParseIntelHex(":020000041234B4\n:00000001FF\n")
	if err != nil {
		t.Error("unexpected error: ", err.Error())
	}
	if m.currentAddress != 0x12340000 {
		t.Errorf("incorrect extended address: %08X", m.currentAddress)
	}
	if len(m.GetDataSegments()) != 0 {
		t.Error("incorrect data segments")
	}
	err = m.ParseIntelHex(":020000049ABCA4\n:00000001FF\n")
	if err != nil {
		t.Error("unexpected error: ", err.Error())
	}
	if m.currentAddress != 0x9ABC0000 {
		t.Errorf("incorrect extended address: %08X", m.currentAddress)
	}
	m.Clear()
	if m.currentAddress != 0 {
		t.Errorf("incorrect extended address: %08X", m.currentAddress)
	}
}

<|endoftext|>


<commit_before>package autofile

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

/*
You can open a Group to keep restrictions on an AutoFile, like
the maximum size of each chunk, and/or the total amount of bytes
stored in the group.

The first file to be written in the Group.Dir is the head file.

  Dir/
  - <HeadPath>

Once the Head file reaches the size limit, it will be rotated.

  Dir/
  - <HeadPath>.000   // First rolled file
  - <HeadPath>       // New head path, starts empty.
                     // The implicit index is 001.

As more files are written, the index numbers grow...

  Dir/
  - <HeadPath>.000   // First rolled file
  - <HeadPath>.001   // Second rolled file
  - ...
  - <HeadPath>       // New head path

The Group can also be used to binary-search for some line,
assuming that marker lines are written occasionally.
*/

const groupCheckDuration = 1000 * time.Millisecond
const defaultHeadSizeLimit = 10 * 1024 * 1024 // 10MB

type Group struct {
	ID             string
	Head           *AutoFile // The head AutoFile to write to
	Dir            string    // Directory that contains .Head
	ticker         *time.Ticker
	mtx            sync.Mutex
	headSizeLimit  int64
	totalSizeLimit int64
	minIndex       int // Includes head
	maxIndex       int // Includes head, where Head will move to
}

func OpenGroup(head *AutoFile) (g *Group, err error) {
	dir := path.Dir(head.Path)

	g = &Group{
		ID:            "group:" + head.ID,
		Head:          head,
		Dir:           dir,
		ticker:        time.NewTicker(groupCheckDuration),
		headSizeLimit: defaultHeadSizeLimit,
		minIndex:      0,
		maxIndex:      0,
	}
	gInfo := g.readGroupInfo()
	g.minIndex = gInfo.MinIndex
	g.maxIndex = gInfo.MaxIndex
	go g.processTicks()
	return
}

func (g *Group) SetHeadSizeLimit(limit int64) {
	g.mtx.Lock()
	g.headSizeLimit = limit
	g.mtx.Unlock()
}

func (g *Group) HeadSizeLimit() int64 {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.headSizeLimit
}

func (g *Group) SetTotalSizeLimit(limit int64) {
	g.mtx.Lock()
	g.totalSizeLimit = limit
	g.mtx.Unlock()
}

func (g *Group) TotalSizeLimit() int64 {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.totalSizeLimit
}

func (g *Group) MaxIndex() int {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.maxIndex
}

func (g *Group) Close() error {
	g.ticker.Stop()
	return nil
}

func (g *Group) processTicks() {
	for {
		_, ok := <-g.ticker.C
		if !ok {
			return // Done.
		}
		g.checkHeadSizeLimit()
		g.checkTotalSizeLimit()
	}
}

// NOTE: for testing
func (g *Group) stopTicker() {
	g.ticker.Stop()
}

// NOTE: this function is called manually in tests.
func (g *Group) checkHeadSizeLimit() {
	size, err := g.Head.Size()
	if err != nil {
		panic(err)
	}
	if size >= g.HeadSizeLimit() {
		g.RotateFile()
	}
}

func (g *Group) checkTotalSizeLimit() {
	// TODO enforce total size limit
}

func (g *Group) RotateFile() {
	g.mtx.Lock()
	defer g.mtx.Unlock()

	dstPath := filePathForIndex(g.Head.Path, g.maxIndex)
	err := os.Rename(g.Head.Path, dstPath)
	if err != nil {
		panic(err)
	}
	err = g.Head.closeFile()
	if err != nil {
		panic(err)
	}
	g.maxIndex += 1
}

func (g *Group) NewReader(index int) *GroupReader {
	r := newGroupReader(g)
	r.SetIndex(index)
	return r
}

// Returns -1 if line comes after, 0 if found, 1 if line comes before.
type SearchFunc func(line string) (int, error)

// Searches for the right file in Group,
// then returns a GroupReader to start streaming lines
// Returns true if an exact match was found, otherwise returns
// the next greater line that starts with prefix.
// CONTRACT: caller is responsible for closing GroupReader.
func (g *Group) Search(prefix string, cmp SearchFunc) (*GroupReader, bool, error) {
	g.mtx.Lock()
	minIndex, maxIndex := g.minIndex, g.maxIndex
	g.mtx.Unlock()
	// Now minIndex/maxIndex may change meanwhile,
	// but it shouldn't be a big deal
	// (maybe we'll want to limit scanUntil though)

	for {
		curIndex := (minIndex + maxIndex + 1) / 2

		// Base case, when there's only 1 choice left.
		if minIndex == maxIndex {
			r := g.NewReader(maxIndex)
			match, err := scanUntil(r, prefix, cmp)
			if err != nil {
				r.Close()
				return nil, false, err
			} else {
				return r, match, err
			}
		}

		// Read starting roughly at the middle file,
		// until we find line that has prefix.
		r := g.NewReader(curIndex)
		foundIndex, line, err := scanFirst(r, prefix)
		r.Close()
		if err != nil {
			return nil, false, err
		}

		// Compare this line to our search query.
		val, err := cmp(line)
		if err != nil {
			return nil, false, err
		}
		if val < 0 {
			// Line will come later
			minIndex = foundIndex
		} else if val == 0 {
			// Stroke of luck, found the line
			r := g.NewReader(foundIndex)
			match, err := scanUntil(r, prefix, cmp)
			if !match {
				panic("Expected match to be true")
			}
			if err != nil {
				r.Close()
				return nil, false, err
			} else {
				return r, true, err
			}
		} else {
			// We passed it
			maxIndex = curIndex - 1
		}
	}

}

// Scans and returns the first line that starts with 'prefix'
func scanFirst(r *GroupReader, prefix string) (int, string, error) {
	for {
		line, err := r.ReadLine()
		if err != nil {
			return 0, "", err
		}
		if !strings.HasPrefix(line, prefix) {
			continue
		}
		index := r.CurIndex()
		return index, line, nil
	}
}

// Returns true iff an exact match was found.
func scanUntil(r *GroupReader, prefix string, cmp SearchFunc) (bool, error) {
	for {
		line, err := r.ReadLine()
		if err != nil {
			return false, err
		}
		if !strings.HasPrefix(line, prefix) {
			continue
		}
		val, err := cmp(line)
		if err != nil {
			return false, err
		}
		if val < 0 {
			continue
		} else if val == 0 {
			r.PushLine(line)
			return true, nil
		} else {
			r.PushLine(line)
			return false, nil
		}
	}
}

type GroupInfo struct {
	MinIndex  int
	MaxIndex  int
	TotalSize int64
	HeadSize  int64
}

// Returns info after scanning all files in g.Head's dir
func (g *Group) ReadGroupInfo() GroupInfo {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.readGroupInfo()
}

// Index includes the head.
// CONTRACT: caller should have called g.mtx.Lock
func (g *Group) readGroupInfo() GroupInfo {
	groupDir := filepath.Dir(g.Head.Path)
	headBase := filepath.Base(g.Head.Path)
	var minIndex, maxIndex int = -1, -1
	var totalSize, headSize int64 = 0, 0

	dir, err := os.Open(groupDir)
	if err != nil {
		panic(err)
	}
	fiz, err := dir.Readdir(0)
	if err != nil {
		panic(err)
	}

	// For each file in the directory, filter by pattern
	for _, fileInfo := range fiz {
		if fileInfo.Name() == headBase {
			fileSize := fileInfo.Size()
			totalSize += fileSize
			headSize = fileSize
			continue
		} else if strings.HasPrefix(fileInfo.Name(), headBase) {
			fileSize := fileInfo.Size()
			totalSize += fileSize
			indexedFilePattern := regexp.MustCompile(`^.+\.([0-9]{3,})$`)
			submatch := indexedFilePattern.FindSubmatch([]byte(fileInfo.Name()))
			if len(submatch) != 0 {
				// Matches
				fileIndex, err := strconv.Atoi(string(submatch[1]))
				if err != nil {
					panic(err)
				}
				if maxIndex < fileIndex {
					maxIndex = fileIndex
				}
				if minIndex == -1 || fileIndex < minIndex {
					minIndex = fileIndex
				}
			}
		}
	}

	// Now account for the head.
	if minIndex == -1 {
		// If there were no numbered files,
		// then the head is index 0.
		minIndex, maxIndex = 0, 0
	} else {
		// Otherwise, the head file is 1 greater
		maxIndex += 1
	}
	return GroupInfo{minIndex, maxIndex, totalSize, headSize}
}

func filePathForIndex(headPath string, index int) string {
	return fmt.Sprintf("%v.%03d", headPath, index)
}

//--------------------------------------------------------------------------------

type GroupReader struct {
	*Group
	mtx       sync.Mutex
	curIndex  int
	curFile   *os.File
	curReader *bufio.Reader
	curLine   []byte
}

func newGroupReader(g *Group) *GroupReader {
	return &GroupReader{
		Group:     g,
		curIndex:  0,
		curFile:   nil,
		curReader: nil,
		curLine:   nil,
	}
}

func (gr *GroupReader) Close() error {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()

	if gr.curReader != nil {
		err := gr.curFile.Close()
		gr.curIndex = 0
		gr.curReader = nil
		gr.curFile = nil
		gr.curLine = nil
		return err
	} else {
		return nil
	}
}

func (gr *GroupReader) ReadLine() (string, error) {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()

	// From PushLine
	if gr.curLine != nil {
		line := string(gr.curLine)
		gr.curLine = nil
		return line, nil
	}

	// Open file if not open yet
	if gr.curReader == nil {
		err := gr.openFile(gr.curIndex)
		if err != nil {
			return "", err
		}
	}

	// Iterate over files until line is found
	for {
		bytes, err := gr.curReader.ReadBytes('\n')
		if err != nil {
			if err != io.EOF {
				return string(bytes), err
			} else {
				// Open the next file
				err := gr.openFile(gr.curIndex + 1)
				if err != nil {
					return "", err
				}
				continue
			}
		}
		return string(bytes), nil
	}
}

// IF index > gr.Group.maxIndex, returns io.EOF
// CONTRACT: caller should hold gr.mtx
func (gr *GroupReader) openFile(index int) error {

	// Lock on Group to ensure that head doesn't move in the meanwhile.
	gr.Group.mtx.Lock()
	defer gr.Group.mtx.Unlock()

	var curFilePath string
	if index == gr.Group.maxIndex {
		curFilePath = gr.Head.Path
	} else if index > gr.Group.maxIndex {
		return io.EOF
	} else {
		curFilePath = filePathForIndex(gr.Head.Path, index)
	}

	curFile, err := os.Open(curFilePath)
	if err != nil {
		return err
	}
	curReader := bufio.NewReader(curFile)

	// Update gr.cur*
	gr.curIndex = index
	gr.curFile = curFile
	gr.curReader = curReader
	gr.curLine = nil
	return nil
}

func (gr *GroupReader) PushLine(line string) {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()

	if gr.curLine == nil {
		gr.curLine = []byte(line)
	} else {
		panic("PushLine failed, already have line")
	}
}

// Cursor's file index.
func (gr *GroupReader) CurIndex() int {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()
	return gr.curIndex
}

func (gr *GroupReader) SetIndex(index int) {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()
	gr.openFile(index)
}
<commit_msg>Add CHALLENGE<commit_after>package autofile

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

/*
You can open a Group to keep restrictions on an AutoFile, like
the maximum size of each chunk, and/or the total amount of bytes
stored in the group.

The first file to be written in the Group.Dir is the head file.

  Dir/
  - <HeadPath>

Once the Head file reaches the size limit, it will be rotated.

  Dir/
  - <HeadPath>.000   // First rolled file
  - <HeadPath>       // New head path, starts empty.
                     // The implicit index is 001.

As more files are written, the index numbers grow...

  Dir/
  - <HeadPath>.000   // First rolled file
  - <HeadPath>.001   // Second rolled file
  - ...
  - <HeadPath>       // New head path

The Group can also be used to binary-search for some line,
assuming that marker lines are written occasionally.
*/

const groupCheckDuration = 1000 * time.Millisecond
const defaultHeadSizeLimit = 10 * 1024 * 1024 // 10MB

type Group struct {
	ID             string
	Head           *AutoFile // The head AutoFile to write to
	Dir            string    // Directory that contains .Head
	ticker         *time.Ticker
	mtx            sync.Mutex
	headSizeLimit  int64
	totalSizeLimit int64
	minIndex       int // Includes head
	maxIndex       int // Includes head, where Head will move to
}

func OpenGroup(head *AutoFile) (g *Group, err error) {
	dir := path.Dir(head.Path)

	g = &Group{
		ID:            "group:" + head.ID,
		Head:          head,
		Dir:           dir,
		ticker:        time.NewTicker(groupCheckDuration),
		headSizeLimit: defaultHeadSizeLimit,
		minIndex:      0,
		maxIndex:      0,
	}
	gInfo := g.readGroupInfo()
	g.minIndex = gInfo.MinIndex
	g.maxIndex = gInfo.MaxIndex
	go g.processTicks()
	return
}

func (g *Group) SetHeadSizeLimit(limit int64) {
	g.mtx.Lock()
	g.headSizeLimit = limit
	g.mtx.Unlock()
}

func (g *Group) HeadSizeLimit() int64 {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.headSizeLimit
}

func (g *Group) SetTotalSizeLimit(limit int64) {
	g.mtx.Lock()
	g.totalSizeLimit = limit
	g.mtx.Unlock()
}

func (g *Group) TotalSizeLimit() int64 {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.totalSizeLimit
}

func (g *Group) MaxIndex() int {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.maxIndex
}

func (g *Group) Close() error {
	g.ticker.Stop()
	return nil
}

func (g *Group) processTicks() {
	for {
		_, ok := <-g.ticker.C
		if !ok {
			return // Done.
		}
		g.checkHeadSizeLimit()
		g.checkTotalSizeLimit()
	}
}

// NOTE: for testing
func (g *Group) stopTicker() {
	g.ticker.Stop()
}

// NOTE: this function is called manually in tests.
func (g *Group) checkHeadSizeLimit() {
	size, err := g.Head.Size()
	if err != nil {
		panic(err)
	}
	if size >= g.HeadSizeLimit() {
		g.RotateFile()
	}
}

func (g *Group) checkTotalSizeLimit() {
	// TODO enforce total size limit
	// CHALLENGE
}

func (g *Group) RotateFile() {
	g.mtx.Lock()
	defer g.mtx.Unlock()

	dstPath := filePathForIndex(g.Head.Path, g.maxIndex)
	err := os.Rename(g.Head.Path, dstPath)
	if err != nil {
		panic(err)
	}
	err = g.Head.closeFile()
	if err != nil {
		panic(err)
	}
	g.maxIndex += 1
}

func (g *Group) NewReader(index int) *GroupReader {
	r := newGroupReader(g)
	r.SetIndex(index)
	return r
}

// Returns -1 if line comes after, 0 if found, 1 if line comes before.
type SearchFunc func(line string) (int, error)

// Searches for the right file in Group,
// then returns a GroupReader to start streaming lines
// Returns true if an exact match was found, otherwise returns
// the next greater line that starts with prefix.
// CONTRACT: caller is responsible for closing GroupReader.
func (g *Group) Search(prefix string, cmp SearchFunc) (*GroupReader, bool, error) {
	g.mtx.Lock()
	minIndex, maxIndex := g.minIndex, g.maxIndex
	g.mtx.Unlock()
	// Now minIndex/maxIndex may change meanwhile,
	// but it shouldn't be a big deal
	// (maybe we'll want to limit scanUntil though)

	for {
		curIndex := (minIndex + maxIndex + 1) / 2

		// Base case, when there's only 1 choice left.
		if minIndex == maxIndex {
			r := g.NewReader(maxIndex)
			match, err := scanUntil(r, prefix, cmp)
			if err != nil {
				r.Close()
				return nil, false, err
			} else {
				return r, match, err
			}
		}

		// Read starting roughly at the middle file,
		// until we find line that has prefix.
		r := g.NewReader(curIndex)
		foundIndex, line, err := scanFirst(r, prefix)
		r.Close()
		if err != nil {
			return nil, false, err
		}

		// Compare this line to our search query.
		val, err := cmp(line)
		if err != nil {
			return nil, false, err
		}
		if val < 0 {
			// Line will come later
			minIndex = foundIndex
		} else if val == 0 {
			// Stroke of luck, found the line
			r := g.NewReader(foundIndex)
			match, err := scanUntil(r, prefix, cmp)
			if !match {
				panic("Expected match to be true")
			}
			if err != nil {
				r.Close()
				return nil, false, err
			} else {
				return r, true, err
			}
		} else {
			// We passed it
			maxIndex = curIndex - 1
		}
	}

}

// Scans and returns the first line that starts with 'prefix'
func scanFirst(r *GroupReader, prefix string) (int, string, error) {
	for {
		line, err := r.ReadLine()
		if err != nil {
			return 0, "", err
		}
		if !strings.HasPrefix(line, prefix) {
			continue
		}
		index := r.CurIndex()
		return index, line, nil
	}
}

// Returns true iff an exact match was found.
func scanUntil(r *GroupReader, prefix string, cmp SearchFunc) (bool, error) {
	for {
		line, err := r.ReadLine()
		if err != nil {
			return false, err
		}
		if !strings.HasPrefix(line, prefix) {
			continue
		}
		val, err := cmp(line)
		if err != nil {
			return false, err
		}
		if val < 0 {
			continue
		} else if val == 0 {
			r.PushLine(line)
			return true, nil
		} else {
			r.PushLine(line)
			return false, nil
		}
	}
}

type GroupInfo struct {
	MinIndex  int
	MaxIndex  int
	TotalSize int64
	HeadSize  int64
}

// Returns info after scanning all files in g.Head's dir
func (g *Group) ReadGroupInfo() GroupInfo {
	g.mtx.Lock()
	defer g.mtx.Unlock()
	return g.readGroupInfo()
}

// Index includes the head.
// CONTRACT: caller should have called g.mtx.Lock
func (g *Group) readGroupInfo() GroupInfo {
	groupDir := filepath.Dir(g.Head.Path)
	headBase := filepath.Base(g.Head.Path)
	var minIndex, maxIndex int = -1, -1
	var totalSize, headSize int64 = 0, 0

	dir, err := os.Open(groupDir)
	if err != nil {
		panic(err)
	}
	fiz, err := dir.Readdir(0)
	if err != nil {
		panic(err)
	}

	// For each file in the directory, filter by pattern
	for _, fileInfo := range fiz {
		if fileInfo.Name() == headBase {
			fileSize := fileInfo.Size()
			totalSize += fileSize
			headSize = fileSize
			continue
		} else if strings.HasPrefix(fileInfo.Name(), headBase) {
			fileSize := fileInfo.Size()
			totalSize += fileSize
			indexedFilePattern := regexp.MustCompile(`^.+\.([0-9]{3,})$`)
			submatch := indexedFilePattern.FindSubmatch([]byte(fileInfo.Name()))
			if len(submatch) != 0 {
				// Matches
				fileIndex, err := strconv.Atoi(string(submatch[1]))
				if err != nil {
					panic(err)
				}
				if maxIndex < fileIndex {
					maxIndex = fileIndex
				}
				if minIndex == -1 || fileIndex < minIndex {
					minIndex = fileIndex
				}
			}
		}
	}

	// Now account for the head.
	if minIndex == -1 {
		// If there were no numbered files,
		// then the head is index 0.
		minIndex, maxIndex = 0, 0
	} else {
		// Otherwise, the head file is 1 greater
		maxIndex += 1
	}
	return GroupInfo{minIndex, maxIndex, totalSize, headSize}
}

func filePathForIndex(headPath string, index int) string {
	return fmt.Sprintf("%v.%03d", headPath, index)
}

//--------------------------------------------------------------------------------

type GroupReader struct {
	*Group
	mtx       sync.Mutex
	curIndex  int
	curFile   *os.File
	curReader *bufio.Reader
	curLine   []byte
}

func newGroupReader(g *Group) *GroupReader {
	return &GroupReader{
		Group:     g,
		curIndex:  0,
		curFile:   nil,
		curReader: nil,
		curLine:   nil,
	}
}

func (gr *GroupReader) Close() error {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()

	if gr.curReader != nil {
		err := gr.curFile.Close()
		gr.curIndex = 0
		gr.curReader = nil
		gr.curFile = nil
		gr.curLine = nil
		return err
	} else {
		return nil
	}
}

func (gr *GroupReader) ReadLine() (string, error) {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()

	// From PushLine
	if gr.curLine != nil {
		line := string(gr.curLine)
		gr.curLine = nil
		return line, nil
	}

	// Open file if not open yet
	if gr.curReader == nil {
		err := gr.openFile(gr.curIndex)
		if err != nil {
			return "", err
		}
	}

	// Iterate over files until line is found
	for {
		bytes, err := gr.curReader.ReadBytes('\n')
		if err != nil {
			if err != io.EOF {
				return string(bytes), err
			} else {
				// Open the next file
				err := gr.openFile(gr.curIndex + 1)
				if err != nil {
					return "", err
				}
				continue
			}
		}
		return string(bytes), nil
	}
}

// IF index > gr.Group.maxIndex, returns io.EOF
// CONTRACT: caller should hold gr.mtx
func (gr *GroupReader) openFile(index int) error {

	// Lock on Group to ensure that head doesn't move in the meanwhile.
	gr.Group.mtx.Lock()
	defer gr.Group.mtx.Unlock()

	var curFilePath string
	if index == gr.Group.maxIndex {
		curFilePath = gr.Head.Path
	} else if index > gr.Group.maxIndex {
		return io.EOF
	} else {
		curFilePath = filePathForIndex(gr.Head.Path, index)
	}

	curFile, err := os.Open(curFilePath)
	if err != nil {
		return err
	}
	curReader := bufio.NewReader(curFile)

	// Update gr.cur*
	gr.curIndex = index
	gr.curFile = curFile
	gr.curReader = curReader
	gr.curLine = nil
	return nil
}

func (gr *GroupReader) PushLine(line string) {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()

	if gr.curLine == nil {
		gr.curLine = []byte(line)
	} else {
		panic("PushLine failed, already have line")
	}
}

// Cursor's file index.
func (gr *GroupReader) CurIndex() int {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()
	return gr.curIndex
}

func (gr *GroupReader) SetIndex(index int) {
	gr.mtx.Lock()
	defer gr.mtx.Unlock()
	gr.openFile(index)
}
<|endoftext|>