<commit_before>/*
 * portal - marshal
 *
 * a library that implements an algorithm for doing consumer coordination within Kafka, rather
 * than using Zookeeper or another external system.
 *
 */

package marshal

import (
	"sync/atomic"
	"time"

	"github.com/jpillora/backoff"
	"github.com/optiopay/kafka"
)

// kafkaConsumerChannel creates a consumer that continuously attempts to consume messages from
// Kafka for the given partition.
func (w *Marshaler) kafkaConsumerChannel(partID int) <-chan message {
	log.Debugf("rationalize[%d]: starting", partID)
	out := make(chan message, 1000)
	go func() {
		var err error
		var alive bool
		var offsetFirst, offsetNext int64

		retry := &backoff.Backoff{Min: 500 * time.Millisecond, Jitter: true}
		for ; true; time.Sleep(retry.Duration()) {
			// Figure out how many messages are in this topic. This can fail if the broker handling
			// this partition is down, so we will loop.
			offsetFirst, err = w.kafka.OffsetEarliest(MarshalTopic, int32(partID))
			if err != nil {
				log.Errorf("rationalize[%d]: failed to get offset: %s", partID, err)
				continue
			}
			offsetNext, err = w.kafka.OffsetLatest(MarshalTopic, int32(partID))
			if err != nil {
				log.Errorf("rationalize[%d]: failed to get offset: %s", partID, err)
				continue
			}
			log.Debugf("rationalize[%d]: offsets %d to %d", partID, offsetFirst, offsetNext)

			// TODO: Is there a case where the latest offset is X>0 but there is no data in
			// the partition? does the offset reset to 0?
			if offsetNext == 0 || offsetFirst == offsetNext {
				alive = true
				w.rationalizers.Done()
			}
			break
		}
		retry.Reset()

		// TODO: Technically we don't have to start at the beginning, we just need to start back
		// a couple heartbeat intervals to get a full state of the world. But this is easiest
		// for right now...
		// TODO: Think about the above. Is it actually true?
		consumerConf := kafka.NewConsumerConf(MarshalTopic, int32(partID))
		consumerConf.StartOffset = kafka.StartOffsetOldest
		consumerConf.RequestTimeout = 1 * time.Second

		consumer, err := w.kafka.Consumer(consumerConf)
		if err != nil {
			// Unfortunately this is a fatal error, as without being able to consume this partition
			// we can't effectively rationalize.
			log.Fatalf("rationalize[%d]: Failed to create consumer: %s", partID, err)
		}

		// Consume messages forever, or until told to quit.
		for {
			if atomic.LoadInt32(w.quit) == 1 {
				log.Debugf("rationalize[%d]: terminating.", partID)
				close(out)
				return
			}

			msgb, err := consumer.Consume()
			if err != nil {
				// The internal consumer will do a number of retries. If we get an error here,
				// we're probably in the middle of a partition handoff. We should pause so we
				// don't hammer the cluster, but otherwise continue.
				log.Warningf("rationalize[%d]: failed to consume: %s", partID, err)
				time.Sleep(retry.Duration())
				continue
			}
			retry.Reset()

			msg, err := decode(msgb.Value)
			if err != nil {
				// Invalid message in the stream. This should never happen, but if it does, just
				// continue on.
				// TODO: We should probably think about this. If we end up in a situation where
				// one version of this software has a bug that writes invalid messages, it could
				// be doing things we don't anticipate. Of course, crashing all consumers
				// reading that partition is also bad.
				log.Errorf("rationalize[%d]: %s", partID, err)
				continue
			}

			log.Debugf("rationalize[%d]: @%d: [%s]", partID, msgb.Offset, msg.Encode())
			out <- msg

			// This is a one-time thing that fires the first time the rationalizer comes up
			// and makes sure we actually process all of the messages.
			if !alive && msgb.Offset >= offsetNext-1 {
				for len(out) > 0 {
					time.Sleep(100 * time.Millisecond)
				}
				log.Infof("rationalize[%d]: reached offset %d, now alive",
					partID, msgb.Offset)
				alive = true
				w.rationalizers.Done()
			}
		}
	}()
	return out
}

// updateClaim is called whenever we need to adjust a claim structure.
func (w *Marshaler) updateClaim(msg *msgHeartbeat) {
	topic := w.getTopicState(msg.GroupID, msg.Topic, msg.PartID)

	topic.lock.Lock()
	defer topic.lock.Unlock()

	// Note that a heartbeat will just set the claim structure. It's not valid to heartbeat
	// for something you don't own (which is why we have ClaimPartition as a separate
	// message), so we can only assume it's valid.
	topic.partitions[msg.PartID].ClientID = msg.ClientID
	topic.partitions[msg.PartID].GroupID = msg.GroupID
	topic.partitions[msg.PartID].LastOffset = msg.LastOffset
	topic.partitions[msg.PartID].LastHeartbeat = int64(msg.Time)
}

// releaseClaim is called whenever someone has released their claim on a partition.
func (w *Marshaler) releaseClaim(msg *msgReleasingPartition) {
	topic := w.getTopicState(msg.GroupID, msg.Topic, msg.PartID)

	topic.lock.Lock()
	defer topic.lock.Unlock()

	// The partition must be claimed by the person releasing it
	if topic.partitions[msg.PartID].ClientID != msg.ClientID ||
		topic.partitions[msg.PartID].GroupID != msg.GroupID {
		log.Warningf("ReleasePartition message from client that doesn't own it. Dropping.")
		return
	}

	// Record the offset they told us they last processed, and then set the heartbeat to 0
	// which means this is no longer claimed
	topic.partitions[msg.PartID].LastOffset = msg.LastOffset
	topic.partitions[msg.PartID].LastHeartbeat = 0
}

// handleClaim is called whenever we see a ClaimPartition message.
func (w *Marshaler) handleClaim(msg *msgClaimingPartition) {
	topic := w.getTopicState(msg.GroupID, msg.Topic, msg.PartID)

	topic.lock.Lock()
	defer topic.lock.Unlock()

	// Send message to all pending consumers then clear the list (it is a violation of the
	// protocol to send two responses)
	fireEvents := func(evt bool) {
		for _, out := range topic.partitions[msg.PartID].pendingClaims {
			out <- evt
		}
		topic.partitions[msg.PartID].pendingClaims = nil
	}

	// Claim logic: if the claim is for an already claimed partition, we can end now and decide
	// whether or not to fire.
	if topic.partitions[msg.PartID].isClaimed(w.ts) {
		// The ClaimPartition message needs to be from us, or we should just return
		if msg.ClientID == w.clientID && msg.GroupID == w.groupID {
			// Now determine if we own the partition claim and let us know whether we do
			// or not
			if topic.partitions[msg.PartID].ClientID == w.clientID &&
				topic.partitions[msg.PartID].GroupID == w.groupID {
				fireEvents(true)
			} else {
				fireEvents(false)
			}
		}
		return
	}

	// At this point, the partition is unclaimed, which means we know we have the first
	// ClaimPartition message. As soon as we get it, we fill in the structure which makes
	// us think it's claimed (it is).
	topic.partitions[msg.PartID].ClientID = msg.ClientID
	topic.partitions[msg.PartID].GroupID = msg.GroupID
	topic.partitions[msg.PartID].LastOffset = 0 // not present in this message, reset.
	topic.partitions[msg.PartID].LastHeartbeat = int64(msg.Time)

	// Now we can advise that this partition has been handled
	if msg.ClientID == w.clientID && msg.GroupID == w.groupID {
		fireEvents(true)
	} else {
		fireEvents(false)
	}
}

// rationalize is a goroutine that constantly consumes from a given partition of the marshal
// topic and makes changes to the world state whenever something happens.
func (w *Marshaler) rationalize(partID int, in <-chan message) { // Might be in over my head.
	for {
		msg, ok := <-in
		if !ok {
			log.Debugf("rationalize[%d]: channel closed.", partID)
			return
		}
		log.Debugf("rationalize[%d]: %s", partID, msg.Encode())

		switch msg.Type() {
		case msgTypeHeartbeat:
			w.updateClaim(msg.(*msgHeartbeat))
		case msgTypeClaimingPartition:
			w.handleClaim(msg.(*msgClaimingPartition))
		case msgTypeReleasingPartition:
			w.releaseClaim(msg.(*msgReleasingPartition))
		case msgTypeClaimingMessages:
			// TODO: Implement.
		}

		// Update step counter so the test suite can wait for messages to be
		// processed in a predictable way (rather than waiting random times)
		atomic.AddInt32(w.rsteps, 1)
	}
}
<commit_msg>Fix hang with first message is an invalid message<commit_after>/*
 * portal - marshal
 *
 * a library that implements an algorithm for doing consumer coordination within Kafka, rather
 * than using Zookeeper or another external system.
 *
 */

package marshal

import (
	"sync/atomic"
	"time"

	"github.com/jpillora/backoff"
	"github.com/optiopay/kafka"
)

// kafkaConsumerChannel creates a consumer that continuously attempts to consume messages from
// Kafka for the given partition.
func (w *Marshaler) kafkaConsumerChannel(partID int) <-chan message {
	log.Debugf("rationalize[%d]: starting", partID)
	out := make(chan message, 1000)
	go func() {
		var err error
		var alive bool
		var offsetFirst, offsetNext int64

		retry := &backoff.Backoff{Min: 500 * time.Millisecond, Jitter: true}
		for ; true; time.Sleep(retry.Duration()) {
			// Figure out how many messages are in this topic. This can fail if the broker handling
			// this partition is down, so we will loop.
			offsetFirst, err = w.kafka.OffsetEarliest(MarshalTopic, int32(partID))
			if err != nil {
				log.Errorf("rationalize[%d]: failed to get offset: %s", partID, err)
				continue
			}
			offsetNext, err = w.kafka.OffsetLatest(MarshalTopic, int32(partID))
			if err != nil {
				log.Errorf("rationalize[%d]: failed to get offset: %s", partID, err)
				continue
			}
			log.Debugf("rationalize[%d]: offsets %d to %d", partID, offsetFirst, offsetNext)

			// TODO: Is there a case where the latest offset is X>0 but there is no data in
			// the partition? does the offset reset to 0?
			if offsetNext == 0 || offsetFirst == offsetNext {
				alive = true
				w.rationalizers.Done()
			}
			break
		}
		retry.Reset()

		// TODO: Technically we don't have to start at the beginning, we just need to start back
		// a couple heartbeat intervals to get a full state of the world. But this is easiest
		// for right now...
		// TODO: Think about the above. Is it actually true?
		consumerConf := kafka.NewConsumerConf(MarshalTopic, int32(partID))
		consumerConf.StartOffset = kafka.StartOffsetOldest
		consumerConf.RequestTimeout = 1 * time.Second

		consumer, err := w.kafka.Consumer(consumerConf)
		if err != nil {
			// Unfortunately this is a fatal error, as without being able to consume this partition
			// we can't effectively rationalize.
			log.Fatalf("rationalize[%d]: Failed to create consumer: %s", partID, err)
		}

		// Consume messages forever, or until told to quit.
		for {
			if atomic.LoadInt32(w.quit) == 1 {
				log.Debugf("rationalize[%d]: terminating.", partID)
				close(out)
				return
			}

			msgb, err := consumer.Consume()
			if err != nil {
				// The internal consumer will do a number of retries. If we get an error here,
				// we're probably in the middle of a partition handoff. We should pause so we
				// don't hammer the cluster, but otherwise continue.
				log.Warningf("rationalize[%d]: failed to consume: %s", partID, err)
				time.Sleep(retry.Duration())
				continue
			}
			retry.Reset()

			msg, err := decode(msgb.Value)
			if err != nil {
				// Invalid message in the stream. This should never happen, but if it does, just
				// continue on.
				// TODO: We should probably think about this. If we end up in a situation where
				// one version of this software has a bug that writes invalid messages, it could
				// be doing things we don't anticipate. Of course, crashing all consumers
				// reading that partition is also bad.
				log.Errorf("rationalize[%d]: %s", partID, err)

				// In the case where the first message is an invalid message, we need to
				// to notify that we're alive now
				if !alive {
					alive = true
					w.rationalizers.Done()
				}
				continue
			}

			log.Debugf("rationalize[%d]: @%d: [%s]", partID, msgb.Offset, msg.Encode())
			out <- msg

			// This is a one-time thing that fires the first time the rationalizer comes up
			// and makes sure we actually process all of the messages.
			if !alive && msgb.Offset >= offsetNext-1 {
				for len(out) > 0 {
					time.Sleep(100 * time.Millisecond)
				}
				log.Infof("rationalize[%d]: reached offset %d, now alive",
					partID, msgb.Offset)
				alive = true
				w.rationalizers.Done()
			}
		}
	}()
	return out
}

// updateClaim is called whenever we need to adjust a claim structure.
func (w *Marshaler) updateClaim(msg *msgHeartbeat) {
	topic := w.getTopicState(msg.GroupID, msg.Topic, msg.PartID)

	topic.lock.Lock()
	defer topic.lock.Unlock()

	// Note that a heartbeat will just set the claim structure. It's not valid to heartbeat
	// for something you don't own (which is why we have ClaimPartition as a separate
	// message), so we can only assume it's valid.
	topic.partitions[msg.PartID].ClientID = msg.ClientID
	topic.partitions[msg.PartID].GroupID = msg.GroupID
	topic.partitions[msg.PartID].LastOffset = msg.LastOffset
	topic.partitions[msg.PartID].LastHeartbeat = int64(msg.Time)
}

// releaseClaim is called whenever someone has released their claim on a partition.
func (w *Marshaler) releaseClaim(msg *msgReleasingPartition) {
	topic := w.getTopicState(msg.GroupID, msg.Topic, msg.PartID)

	topic.lock.Lock()
	defer topic.lock.Unlock()

	// The partition must be claimed by the person releasing it
	if topic.partitions[msg.PartID].ClientID != msg.ClientID ||
		topic.partitions[msg.PartID].GroupID != msg.GroupID {
		log.Warningf("ReleasePartition message from client that doesn't own it. Dropping.")
		return
	}

	// Record the offset they told us they last processed, and then set the heartbeat to 0
	// which means this is no longer claimed
	topic.partitions[msg.PartID].LastOffset = msg.LastOffset
	topic.partitions[msg.PartID].LastHeartbeat = 0
}

// handleClaim is called whenever we see a ClaimPartition message.
func (w *Marshaler) handleClaim(msg *msgClaimingPartition) {
	topic := w.getTopicState(msg.GroupID, msg.Topic, msg.PartID)

	topic.lock.Lock()
	defer topic.lock.Unlock()

	// Send message to all pending consumers then clear the list (it is a violation of the
	// protocol to send two responses)
	fireEvents := func(evt bool) {
		for _, out := range topic.partitions[msg.PartID].pendingClaims {
			out <- evt
		}
		topic.partitions[msg.PartID].pendingClaims = nil
	}

	// Claim logic: if the claim is for an already claimed partition, we can end now and decide
	// whether or not to fire.
	if topic.partitions[msg.PartID].isClaimed(w.ts) {
		// The ClaimPartition message needs to be from us, or we should just return
		if msg.ClientID == w.clientID && msg.GroupID == w.groupID {
			// Now determine if we own the partition claim and let us know whether we do
			// or not
			if topic.partitions[msg.PartID].ClientID == w.clientID &&
				topic.partitions[msg.PartID].GroupID == w.groupID {
				fireEvents(true)
			} else {
				fireEvents(false)
			}
		}
		return
	}

	// At this point, the partition is unclaimed, which means we know we have the first
	// ClaimPartition message. As soon as we get it, we fill in the structure which makes
	// us think it's claimed (it is).
	topic.partitions[msg.PartID].ClientID = msg.ClientID
	topic.partitions[msg.PartID].GroupID = msg.GroupID
	topic.partitions[msg.PartID].LastOffset = 0 // not present in this message, reset.
	topic.partitions[msg.PartID].LastHeartbeat = int64(msg.Time)

	// Now we can advise that this partition has been handled
	if msg.ClientID == w.clientID && msg.GroupID == w.groupID {
		fireEvents(true)
	} else {
		fireEvents(false)
	}
}

// rationalize is a goroutine that constantly consumes from a given partition of the marshal
// topic and makes changes to the world state whenever something happens.
func (w *Marshaler) rationalize(partID int, in <-chan message) { // Might be in over my head.
	for {
		msg, ok := <-in
		if !ok {
			log.Debugf("rationalize[%d]: channel closed.", partID)
			return
		}
		log.Debugf("rationalize[%d]: %s", partID, msg.Encode())

		switch msg.Type() {
		case msgTypeHeartbeat:
			w.updateClaim(msg.(*msgHeartbeat))
		case msgTypeClaimingPartition:
			w.handleClaim(msg.(*msgClaimingPartition))
		case msgTypeReleasingPartition:
			w.releaseClaim(msg.(*msgReleasingPartition))
		case msgTypeClaimingMessages:
			// TODO: Implement.
		}

		// Update step counter so the test suite can wait for messages to be
		// processed in a predictable way (rather than waiting random times)
		atomic.AddInt32(w.rsteps, 1)
	}
}
<|endoftext|>


<commit_before>/*
 *    Copyright (C) 2014 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

// beehive's web-module.
package webbee

import (
	"encoding/json"
	"github.com/hoisie/web"
	"github.com/muesli/beehive/modules"
	"io/ioutil"
	"log"
)

type WebBee struct {
	modules.Module

	addr        string
	path		string

	eventChan chan modules.Event
}

func (mod *WebBee) Run(cin chan modules.Event) {
	mod.eventChan = cin

	web.Get(mod.path, mod.GetRequest)
	web.Post(mod.path, mod.PostRequest)

	web.Run(mod.addr)
}

func (mod *WebBee) Action(action modules.Action) []modules.Placeholder {
	outs := []modules.Placeholder{}

	switch action.Name {
	case "post":
/*		json := ""

		for _, opt := range action.Options {
			if opt.Name == "json" {
				json = opt.Value.(string)
			}
		}*/

		//FIXME: do stuff

	default:
		panic("Unknown action triggered in " +mod.Name()+": "+action.Name)
	}

	return outs
}

func (mod *WebBee) GetRequest(ctx *web.Context) {
	ev := modules.Event{
		Bee:  mod.Name(),
		Name: "get",
		Options: []modules.Placeholder{
			modules.Placeholder{
				Name:  "ip",
				Type:  "string",
				Value: "tbd",
			},
		},
	}

	for k, v := range ctx.Params {
		log.Println("GET param:", k, "=", v)

		ph := modules.Placeholder{
			Name: k,
			Type: "string",
			Value: v,
		}
		ev.Options = append(ev.Options, ph)
	}

	mod.eventChan <- ev
}

func (mod *WebBee) PostRequest(ctx *web.Context) {
	b, err := ioutil.ReadAll(ctx.Request.Body)
	if err != nil {
		log.Println("Error:", err)
		return
	}

	var payload interface{}
	err = json.Unmarshal(b, &payload)
	if err != nil {
		log.Println("Error:", err)
		return
	}

	ev := modules.Event{
		Bee:  mod.Name(),
		Name: "post",
		Options: []modules.Placeholder{
			modules.Placeholder{
				Name:  "json",
				Type:  "map",
				Value: payload,
			},
			modules.Placeholder{
				Name:  "ip",
				Type:  "string",
				Value: "tbd",
			},
		},
	}

	j := make(map[string]interface{})
	body, err := ioutil.ReadAll(ctx.Request.Body)
	if err != nil {
		log.Println("Error:", err)
	  return
	}
	err = json.Unmarshal(body, &j)
	if err != nil {
		log.Println("Error:", err)
	  return
	}

	for k, v := range j {
		log.Printf("POST JSON param: %s = %+v\n", k, v)

		ph := modules.Placeholder{
			Name: k,
			Type: "string",
			Value: v,
		}
		ev.Options = append(ev.Options, ph)
	}

	mod.eventChan <- ev
}
<commit_msg>* Body has already been read.<commit_after>/*
 *    Copyright (C) 2014 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

// beehive's web-module.
package webbee

import (
	"encoding/json"
	"github.com/hoisie/web"
	"github.com/muesli/beehive/modules"
	"io/ioutil"
	"log"
)

type WebBee struct {
	modules.Module

	addr        string
	path		string

	eventChan chan modules.Event
}

func (mod *WebBee) Run(cin chan modules.Event) {
	mod.eventChan = cin

	web.Get(mod.path, mod.GetRequest)
	web.Post(mod.path, mod.PostRequest)

	web.Run(mod.addr)
}

func (mod *WebBee) Action(action modules.Action) []modules.Placeholder {
	outs := []modules.Placeholder{}

	switch action.Name {
	case "post":
/*		json := ""

		for _, opt := range action.Options {
			if opt.Name == "json" {
				json = opt.Value.(string)
			}
		}*/

		//FIXME: do stuff

	default:
		panic("Unknown action triggered in " +mod.Name()+": "+action.Name)
	}

	return outs
}

func (mod *WebBee) GetRequest(ctx *web.Context) {
	ev := modules.Event{
		Bee:  mod.Name(),
		Name: "get",
		Options: []modules.Placeholder{
			modules.Placeholder{
				Name:  "ip",
				Type:  "string",
				Value: "tbd",
			},
		},
	}

	for k, v := range ctx.Params {
		log.Println("GET param:", k, "=", v)

		ph := modules.Placeholder{
			Name: k,
			Type: "string",
			Value: v,
		}
		ev.Options = append(ev.Options, ph)
	}

	mod.eventChan <- ev
}

func (mod *WebBee) PostRequest(ctx *web.Context) {
	b, err := ioutil.ReadAll(ctx.Request.Body)
	if err != nil {
		log.Println("Error:", err)
		return
	}

	var payload interface{}
	err = json.Unmarshal(b, &payload)
	if err != nil {
		log.Println("Error:", err)
		return
	}

	ev := modules.Event{
		Bee:  mod.Name(),
		Name: "post",
		Options: []modules.Placeholder{
			modules.Placeholder{
				Name:  "json",
				Type:  "map",
				Value: payload,
			},
			modules.Placeholder{
				Name:  "ip",
				Type:  "string",
				Value: "tbd",
			},
		},
	}

	j := make(map[string]interface{})
	err = json.Unmarshal(b, &j)
	if err != nil {
		log.Println("Error:", err)
	  return
	}

	for k, v := range j {
		log.Printf("POST JSON param: %s = %+v\n", k, v)

		ph := modules.Placeholder{
			Name: k,
			Type: "string",
			Value: v,
		}
		ev.Options = append(ev.Options, ph)
	}

	mod.eventChan <- ev
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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 resourcequota

import (
	"fmt"
	"io"
	"strings"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/admission"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/resource"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/resourcequota"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
)

func init() {
	admission.RegisterPlugin("ResourceQuota", func(client client.Interface, config io.Reader) (admission.Interface, error) {
		return NewResourceQuota(client), nil
	})
}

type quota struct {
	client  client.Interface
	indexer cache.Indexer
}

func NewResourceQuota(client client.Interface) admission.Interface {
	lw := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return client.ResourceQuotas(api.NamespaceAll).List(labels.Everything())
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return client.ResourceQuotas(api.NamespaceAll).Watch(labels.Everything(), fields.Everything(), resourceVersion)
		},
	}
	indexer, reflector := cache.NewNamespaceKeyedIndexerAndReflector(lw, &api.ResourceQuota{}, 0)
	reflector.Run()
	return &quota{client: client, indexer: indexer}
}

var resourceToResourceName = map[string]api.ResourceName{
	"pods":                   api.ResourcePods,
	"services":               api.ResourceServices,
	"replicationcontrollers": api.ResourceReplicationControllers,
	"resourcequotas":         api.ResourceQuotas,
	"secrets":                api.ResourceSecrets,
	"persistentvolumeclaims": api.ResourcePersistentVolumeClaims,
}

func (q *quota) Admit(a admission.Attributes) (err error) {
	if a.GetOperation() == "DELETE" {
		return nil
	}

	key := &api.ResourceQuota{
		ObjectMeta: api.ObjectMeta{
			Namespace: a.GetNamespace(),
			Name:      "",
		},
	}
	items, err := q.indexer.Index("namespace", key)
	if err != nil {
		return admission.NewForbidden(a, fmt.Errorf("Unable to %s %s at this time because there was an error enforcing quota", a.GetOperation(), a.GetResource()))
	}
	if len(items) == 0 {
		return nil
	}

	for i := range items {
		quota := items[i].(*api.ResourceQuota)

		// we cannot modify the value directly in the cache, so we copy
		status := &api.ResourceQuotaStatus{
			Hard: api.ResourceList{},
			Used: api.ResourceList{},
		}
		for k, v := range quota.Status.Hard {
			status.Hard[k] = *v.Copy()
		}
		for k, v := range quota.Status.Used {
			status.Used[k] = *v.Copy()
		}

		dirty, err := IncrementUsage(a, status, q.client)
		if err != nil {
			return admission.NewForbidden(a, err)
		}

		if dirty {
			// construct a usage record
			usage := api.ResourceQuota{
				ObjectMeta: api.ObjectMeta{
					Name:            quota.Name,
					Namespace:       quota.Namespace,
					ResourceVersion: quota.ResourceVersion,
					Labels:          quota.Labels,
					Annotations:     quota.Annotations},
			}
			usage.Status = *status
			_, err = q.client.ResourceQuotas(usage.Namespace).UpdateStatus(&usage)
			if err != nil {
				return admission.NewForbidden(a, fmt.Errorf("Unable to %s %s at this time because there was an error enforcing quota", a.GetOperation(), a.GetResource()))
			}
		}
	}
	return nil
}

// IncrementUsage updates the supplied ResourceQuotaStatus object based on the incoming operation
// Return true if the usage must be recorded prior to admitting the new resource
// Return an error if the operation should not pass admission control
func IncrementUsage(a admission.Attributes, status *api.ResourceQuotaStatus, client client.Interface) (bool, error) {
	dirty := false
	set := map[api.ResourceName]bool{}
	for k := range status.Hard {
		set[k] = true
	}
	obj := a.GetObject()
	// handle max counts for each kind of resource (pods, services, replicationControllers, etc.)
	if a.GetOperation() == "CREATE" {
		// TODO v1beta1 had camel case, v1beta3 went to all lower, we can remove this line when we deprecate v1beta1
		resourceNormalized := strings.ToLower(a.GetResource())
		resourceName := resourceToResourceName[resourceNormalized]
		hard, hardFound := status.Hard[resourceName]
		if hardFound {
			used, usedFound := status.Used[resourceName]
			if !usedFound {
				return false, fmt.Errorf("Quota usage stats are not yet known, unable to admit resource until an accurate count is completed.")
			}
			if used.Value() >= hard.Value() {
				return false, fmt.Errorf("Limited to %s %s", hard.String(), resourceName)
			} else {
				status.Used[resourceName] = *resource.NewQuantity(used.Value()+int64(1), resource.DecimalSI)
				dirty = true
			}
		}
	}
	// handle memory/cpu constraints, and any diff of usage based on memory/cpu on updates
	if a.GetResource() == "pods" && (set[api.ResourceMemory] || set[api.ResourceCPU]) {
		pod := obj.(*api.Pod)
		deltaCPU := resourcequota.PodCPU(pod)
		deltaMemory := resourcequota.PodMemory(pod)
		// if this is an update, we need to find the delta cpu/memory usage from previous state
		if a.GetOperation() == "UPDATE" {
			oldPod, err := client.Pods(a.GetNamespace()).Get(pod.Name)
			if err != nil {
				return false, err
			}
			oldCPU := resourcequota.PodCPU(oldPod)
			oldMemory := resourcequota.PodMemory(oldPod)
			deltaCPU = resource.NewMilliQuantity(deltaCPU.MilliValue()-oldCPU.MilliValue(), resource.DecimalSI)
			deltaMemory = resource.NewQuantity(deltaMemory.Value()-oldMemory.Value(), resource.DecimalSI)
		}

		hardMem, hardMemFound := status.Hard[api.ResourceMemory]
		if hardMemFound {
			if set[api.ResourceMemory] && resourcequota.IsPodMemoryUnbounded(pod) {
				return false, fmt.Errorf("Limited to %s memory, but pod has no specified memory limit", hardMem.String())
			}
			used, usedFound := status.Used[api.ResourceMemory]
			if !usedFound {
				return false, fmt.Errorf("Quota usage stats are not yet known, unable to admit resource until an accurate count is completed.")
			}
			if used.Value()+deltaMemory.Value() > hardMem.Value() {
				return false, fmt.Errorf("Limited to %s memory", hardMem.String())
			} else {
				status.Used[api.ResourceMemory] = *resource.NewQuantity(used.Value()+deltaMemory.Value(), resource.DecimalSI)
				dirty = true
			}
		}
		hardCPU, hardCPUFound := status.Hard[api.ResourceCPU]
		if hardCPUFound {
			if set[api.ResourceCPU] && resourcequota.IsPodCPUUnbounded(pod) {
				return false, fmt.Errorf("Limited to %s CPU, but pod has no specified cpu limit", hardCPU.String())
			}
			used, usedFound := status.Used[api.ResourceCPU]
			if !usedFound {
				return false, fmt.Errorf("Quota usage stats are not yet known, unable to admit resource until an accurate count is completed.")
			}
			if used.MilliValue()+deltaCPU.MilliValue() > hardCPU.MilliValue() {
				return false, fmt.Errorf("Limited to %s CPU", hardCPU.String())
			} else {
				status.Used[api.ResourceCPU] = *resource.NewMilliQuantity(used.MilliValue()+deltaCPU.MilliValue(), resource.DecimalSI)
				dirty = true
			}
		}
	}
	return dirty, nil
}
<commit_msg>Retry incrementing quota if there is a conflict<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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 resourcequota

import (
	"fmt"
	"io"
	"math/rand"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/admission"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/resource"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/resourcequota"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
)

func init() {
	admission.RegisterPlugin("ResourceQuota", func(client client.Interface, config io.Reader) (admission.Interface, error) {
		return NewResourceQuota(client), nil
	})
}

type quota struct {
	client  client.Interface
	indexer cache.Indexer
}

func NewResourceQuota(client client.Interface) admission.Interface {
	lw := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return client.ResourceQuotas(api.NamespaceAll).List(labels.Everything())
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return client.ResourceQuotas(api.NamespaceAll).Watch(labels.Everything(), fields.Everything(), resourceVersion)
		},
	}
	indexer, reflector := cache.NewNamespaceKeyedIndexerAndReflector(lw, &api.ResourceQuota{}, 0)
	reflector.Run()
	return &quota{client: client, indexer: indexer}
}

var resourceToResourceName = map[string]api.ResourceName{
	"pods":                   api.ResourcePods,
	"services":               api.ResourceServices,
	"replicationcontrollers": api.ResourceReplicationControllers,
	"resourcequotas":         api.ResourceQuotas,
	"secrets":                api.ResourceSecrets,
	"persistentvolumeclaims": api.ResourcePersistentVolumeClaims,
}

func (q *quota) Admit(a admission.Attributes) (err error) {
	if a.GetOperation() == "DELETE" {
		return nil
	}

	key := &api.ResourceQuota{
		ObjectMeta: api.ObjectMeta{
			Namespace: a.GetNamespace(),
			Name:      "",
		},
	}

	// concurrent operations that modify quota tracked resources can cause a conflict when incrementing usage
	// as a result, we will attempt to increment quota usage per request up to numRetries limit
	// we fuzz each retry with an interval period to attempt to improve end-user experience during concurrent operations
	numRetries := 10
	interval := time.Duration(rand.Int63n(90)+int64(10)) * time.Millisecond

	items, err := q.indexer.Index("namespace", key)
	if err != nil {
		return admission.NewForbidden(a, fmt.Errorf("Unable to %s %s at this time because there was an error enforcing quota", a.GetOperation(), a.GetResource()))
	}
	if len(items) == 0 {
		return nil
	}

	for i := range items {

		quota := items[i].(*api.ResourceQuota)

		for retry := 1; retry <= numRetries; retry++ {

			// we cannot modify the value directly in the cache, so we copy
			status := &api.ResourceQuotaStatus{
				Hard: api.ResourceList{},
				Used: api.ResourceList{},
			}
			for k, v := range quota.Status.Hard {
				status.Hard[k] = *v.Copy()
			}
			for k, v := range quota.Status.Used {
				status.Used[k] = *v.Copy()
			}

			dirty, err := IncrementUsage(a, status, q.client)
			if err != nil {
				return admission.NewForbidden(a, err)
			}

			if dirty {
				// construct a usage record
				usage := api.ResourceQuota{
					ObjectMeta: api.ObjectMeta{
						Name:            quota.Name,
						Namespace:       quota.Namespace,
						ResourceVersion: quota.ResourceVersion,
						Labels:          quota.Labels,
						Annotations:     quota.Annotations},
				}
				usage.Status = *status
				_, err = q.client.ResourceQuotas(usage.Namespace).UpdateStatus(&usage)
				if err == nil {
					break
				}

				// we have concurrent requests to update quota, so look to retry if needed
				if retry == numRetries {
					return admission.NewForbidden(a, fmt.Errorf("Unable to %s %s at this time because there are too many concurrent requests to increment quota", a.GetOperation(), a.GetResource()))
				}
				time.Sleep(interval)
				// manually get the latest quota
				quota, err = q.client.ResourceQuotas(usage.Namespace).Get(quota.Name)
				if err != nil {
					return admission.NewForbidden(a, err)
				}
			}
		}
	}
	return nil
}

// IncrementUsage updates the supplied ResourceQuotaStatus object based on the incoming operation
// Return true if the usage must be recorded prior to admitting the new resource
// Return an error if the operation should not pass admission control
func IncrementUsage(a admission.Attributes, status *api.ResourceQuotaStatus, client client.Interface) (bool, error) {
	dirty := false
	set := map[api.ResourceName]bool{}
	for k := range status.Hard {
		set[k] = true
	}
	obj := a.GetObject()
	// handle max counts for each kind of resource (pods, services, replicationControllers, etc.)
	if a.GetOperation() == "CREATE" {
		// TODO v1beta1 had camel case, v1beta3 went to all lower, we can remove this line when we deprecate v1beta1
		resourceNormalized := strings.ToLower(a.GetResource())
		resourceName := resourceToResourceName[resourceNormalized]
		hard, hardFound := status.Hard[resourceName]
		if hardFound {
			used, usedFound := status.Used[resourceName]
			if !usedFound {
				return false, fmt.Errorf("Quota usage stats are not yet known, unable to admit resource until an accurate count is completed.")
			}
			if used.Value() >= hard.Value() {
				return false, fmt.Errorf("Limited to %s %s", hard.String(), resourceName)
			} else {
				status.Used[resourceName] = *resource.NewQuantity(used.Value()+int64(1), resource.DecimalSI)
				dirty = true
			}
		}
	}
	// handle memory/cpu constraints, and any diff of usage based on memory/cpu on updates
	if a.GetResource() == "pods" && (set[api.ResourceMemory] || set[api.ResourceCPU]) {
		pod := obj.(*api.Pod)
		deltaCPU := resourcequota.PodCPU(pod)
		deltaMemory := resourcequota.PodMemory(pod)
		// if this is an update, we need to find the delta cpu/memory usage from previous state
		if a.GetOperation() == "UPDATE" {
			oldPod, err := client.Pods(a.GetNamespace()).Get(pod.Name)
			if err != nil {
				return false, err
			}
			oldCPU := resourcequota.PodCPU(oldPod)
			oldMemory := resourcequota.PodMemory(oldPod)
			deltaCPU = resource.NewMilliQuantity(deltaCPU.MilliValue()-oldCPU.MilliValue(), resource.DecimalSI)
			deltaMemory = resource.NewQuantity(deltaMemory.Value()-oldMemory.Value(), resource.DecimalSI)
		}

		hardMem, hardMemFound := status.Hard[api.ResourceMemory]
		if hardMemFound {
			if set[api.ResourceMemory] && resourcequota.IsPodMemoryUnbounded(pod) {
				return false, fmt.Errorf("Limited to %s memory, but pod has no specified memory limit", hardMem.String())
			}
			used, usedFound := status.Used[api.ResourceMemory]
			if !usedFound {
				return false, fmt.Errorf("Quota usage stats are not yet known, unable to admit resource until an accurate count is completed.")
			}
			if used.Value()+deltaMemory.Value() > hardMem.Value() {
				return false, fmt.Errorf("Limited to %s memory", hardMem.String())
			} else {
				status.Used[api.ResourceMemory] = *resource.NewQuantity(used.Value()+deltaMemory.Value(), resource.DecimalSI)
				dirty = true
			}
		}
		hardCPU, hardCPUFound := status.Hard[api.ResourceCPU]
		if hardCPUFound {
			if set[api.ResourceCPU] && resourcequota.IsPodCPUUnbounded(pod) {
				return false, fmt.Errorf("Limited to %s CPU, but pod has no specified cpu limit", hardCPU.String())
			}
			used, usedFound := status.Used[api.ResourceCPU]
			if !usedFound {
				return false, fmt.Errorf("Quota usage stats are not yet known, unable to admit resource until an accurate count is completed.")
			}
			if used.MilliValue()+deltaCPU.MilliValue() > hardCPU.MilliValue() {
				return false, fmt.Errorf("Limited to %s CPU", hardCPU.String())
			} else {
				status.Used[api.ResourceCPU] = *resource.NewMilliQuantity(used.MilliValue()+deltaCPU.MilliValue(), resource.DecimalSI)
				dirty = true
			}
		}
	}
	return dirty, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"github.com/darkhelmet/twitterstream"
	"os"
	"sync"
)

type Secrets struct {
	ConsumerKey    string `json:"consumerKey,omitempty"`
	ConsumerSecret string `json:"consumerSecret,omitempty"`
	AccessToken    string `json:"accessToken,omitempty"`
	AccessSecret   string `json:"accessSecret,omitempty"`
}

type TwitterPlugin struct {
	sync.RWMutex
	client     *twitterstream.Client
	conn       *twitterstream.Connection
	config     *Secrets
	Subscribed map[*Client]byte
	stop       chan bool
}

func init() {
	pluginMutex.Lock()
	if pluginNodes == nil {
		pluginNodes = make(map[string]Plugin)
	}
	pluginNodes["$twitter"] = &TwitterPlugin{}
	pluginMutex.Unlock()
}

func (tp *TwitterPlugin) Initialise() error {
	if err := tp.ReadConfig(); err != nil {
		return err
	}
	if tp.config.ConsumerKey == "" || tp.config.ConsumerSecret == "" || tp.config.AccessToken == "" || tp.config.AccessSecret == "" {
		return errors.New("Not all twitter secrets defined")
	}
	tp.Subscribed = make(map[*Client]byte)
	tp.client = twitterstream.NewClient(tp.config.ConsumerKey, tp.config.ConsumerSecret, tp.config.AccessToken, tp.config.AccessSecret)
	return nil
}

func (tp *TwitterPlugin) AddSub(client *Client, subscription []string, qos byte, complete chan bool) {
	tp.Lock()
	defer func() {
		complete <- true
		tp.Subscribed[client] = qos
		tp.Unlock()
	}()
	var err error
	INFO.Println("Adding $twitter sub for", subscription[1], client.clientId)
	if tp.conn != nil {
		close(tp.stop)
		tp.conn.Close()
	}
	tp.stop = make(chan bool)
	tp.conn, err = tp.client.Track(subscription[1])
	if err != nil {
		ERROR.Println(err.Error())
		return
	}
	go tp.Run()
}

func (tp *TwitterPlugin) DeleteSub(client *Client, complete chan bool) {
	tp.Lock()
	defer tp.Unlock()
	delete(tp.Subscribed, client)
	if len(tp.Subscribed) == 0 && tp.conn != nil {
		INFO.Println("All subscriptions gone, closing twitter connection")
		close(tp.stop)
		tp.conn.Close()
	}
	complete <- true
}

func (tp *TwitterPlugin) ReadConfig() error {
	file, err := os.Open("twitter_plugin_config.json")
	if err != nil {
		return err
	}
	decoder := json.NewDecoder(file)

	err = decoder.Decode(&tp.config)
	if err != nil {
		return err
	}
	return nil
}

func (tp *TwitterPlugin) Run() {
	tweetChan := make(chan *twitterstream.Tweet, 2)
	go func() {
		for {
			tweet, err := tp.conn.Next()
			if err != nil {
				ERROR.Println("Twitter receive error")
				return
			}
			tweetChan <- tweet
			select {
			case <-tp.stop:
				return
			default:
			}
		}
	}()

	for {
		select {
		case tweet := <-tweetChan:
			tp.RLock()
			message := New(PUBLISH).(*publishPacket)
			message.Qos = 0
			message.topicName = "$twitter/" + tweet.User.ScreenName
			message.payload = []byte(tweet.Text)
			for client, _ := range tp.Subscribed {
				client.outboundMessages.Push(message)
			}
			tp.RUnlock()
		case <-tp.stop:
			return
		}
	}
}
<commit_msg>reset plugin connection when stopped<commit_after>package main

import (
	"encoding/json"
	"errors"
	"github.com/darkhelmet/twitterstream"
	"os"
	"sync"
)

type Secrets struct {
	ConsumerKey    string `json:"consumerKey,omitempty"`
	ConsumerSecret string `json:"consumerSecret,omitempty"`
	AccessToken    string `json:"accessToken,omitempty"`
	AccessSecret   string `json:"accessSecret,omitempty"`
}

type TwitterPlugin struct {
	sync.RWMutex
	client     *twitterstream.Client
	conn       *twitterstream.Connection
	config     *Secrets
	Subscribed map[*Client]byte
	stop       chan bool
}

func init() {
	pluginMutex.Lock()
	if pluginNodes == nil {
		pluginNodes = make(map[string]Plugin)
	}
	pluginNodes["$twitter"] = &TwitterPlugin{}
	pluginMutex.Unlock()
}

func (tp *TwitterPlugin) Initialise() error {
	if err := tp.ReadConfig(); err != nil {
		return err
	}
	if tp.config.ConsumerKey == "" || tp.config.ConsumerSecret == "" || tp.config.AccessToken == "" || tp.config.AccessSecret == "" {
		return errors.New("Not all twitter secrets defined")
	}
	tp.Subscribed = make(map[*Client]byte)
	tp.client = twitterstream.NewClient(tp.config.ConsumerKey, tp.config.ConsumerSecret, tp.config.AccessToken, tp.config.AccessSecret)
	return nil
}

func (tp *TwitterPlugin) AddSub(client *Client, subscription []string, qos byte, complete chan bool) {
	tp.Lock()
	defer func() {
		complete <- true
		tp.Subscribed[client] = qos
		tp.Unlock()
	}()
	var err error
	INFO.Println("Adding $twitter sub for", subscription[1], client.clientId)
	if tp.conn != nil {
		close(tp.stop)
		tp.conn.Close()
	}
	tp.stop = make(chan bool)
	tp.conn, err = tp.client.Track(subscription[1])
	if err != nil {
		ERROR.Println(err.Error())
		return
	}
	go tp.Run()
}

func (tp *TwitterPlugin) DeleteSub(client *Client, complete chan bool) {
	tp.Lock()
	defer tp.Unlock()
	delete(tp.Subscribed, client)
	if len(tp.Subscribed) == 0 && tp.conn != nil {
		INFO.Println("All subscriptions gone, closing twitter connection")
		close(tp.stop)
		tp.conn.Close()
		tp.conn == nil
	}
	complete <- true
}

func (tp *TwitterPlugin) ReadConfig() error {
	file, err := os.Open("twitter_plugin_config.json")
	if err != nil {
		return err
	}
	decoder := json.NewDecoder(file)

	err = decoder.Decode(&tp.config)
	if err != nil {
		return err
	}
	return nil
}

func (tp *TwitterPlugin) Run() {
	tweetChan := make(chan *twitterstream.Tweet, 2)
	go func() {
		for {
			tweet, err := tp.conn.Next()
			if err != nil {
				ERROR.Println("Twitter receive error")
				return
			}
			tweetChan <- tweet
			select {
			case <-tp.stop:
				return
			default:
			}
		}
	}()

	for {
		select {
		case tweet := <-tweetChan:
			tp.RLock()
			message := New(PUBLISH).(*publishPacket)
			message.Qos = 0
			message.topicName = "$twitter/" + tweet.User.ScreenName
			message.payload = []byte(tweet.Text)
			for client, _ := range tp.Subscribed {
				client.outboundMessages.Push(message)
			}
			tp.RUnlock()
		case <-tp.stop:
			return
		}
	}
}
<|endoftext|>


<commit_before>// +build linux darwin freebsd netbsd openbsd solaris dragonfly windows

package pb

import (
	"io"
	"sync"
	"time"
)

// Create and start new pool with given bars
// You need call pool.Stop() after work
func StartPool(pbs ...*ProgressBar) (pool *Pool, err error) {
	pool = new(Pool)
	if err = pool.start(); err != nil {
		return
	}
	pool.Add(pbs...)
	return
}

type Pool struct {
	Output        io.Writer
	RefreshRate   time.Duration
	bars          []*ProgressBar
	lastBarsCount int
	shutdownCh    chan struct{}
	workerCh      chan struct{}
	m             sync.Mutex
	finishOnce    sync.Once
}

// Add progress bars.
func (p *Pool) Add(pbs ...*ProgressBar) {
	p.m.Lock()
	defer p.m.Unlock()
	for _, bar := range pbs {
		bar.ManualUpdate = true
		bar.NotPrint = true
		bar.Start()
		p.bars = append(p.bars, bar)
	}
}

func (p *Pool) start() (err error) {
	p.RefreshRate = DefaultRefreshRate
	p.shutdownCh, err = lockEcho()
	if err != nil {
		return
	}
	p.workerCh = make(chan struct{})
	go p.writer()
	return
}

func (p *Pool) writer() {
	var first = true
	defer func() {
		if first == false {
			p.print(false)
		} else {
			p.print(true)
			p.print(false)
		}
		close(p.workerCh)
	}()

	for {
		select {
		case <-time.After(p.RefreshRate):
			if p.print(first) {
				p.print(false)
				return
			}
			first = false
		case <-p.shutdownCh:
			return
		}
	}
}

// Restore terminal state and close pool
func (p *Pool) Stop() error {
	p.finishOnce.Do(func() {
		close(p.shutdownCh)
	})

	// Wait for the worker to complete
	select {
	case <-p.workerCh:
	}

	return unlockEcho()
}
<commit_msg>Allow for initialising a pool without starting it<commit_after>// +build linux darwin freebsd netbsd openbsd solaris dragonfly windows

package pb

import (
	"io"
	"sync"
	"time"
)

// Create and start new pool with given bars
// You need call pool.Stop() after work
func StartPool(pbs ...*ProgressBar) (pool *Pool, err error) {
	pool = new(Pool)
	if err = pool.Start(); err != nil {
		return
	}
	pool.Add(pbs...)
	return
}

// NewPool initialises a pool with progress bars, but
// doesn't start it. You need to call Start manually
func NewPool(pbs ...*ProgressBar) (pool *Pool) {
	pool = new(Pool)
	pool.Add(pbs...)
	return
}

type Pool struct {
	Output        io.Writer
	RefreshRate   time.Duration
	bars          []*ProgressBar
	lastBarsCount int
	shutdownCh    chan struct{}
	workerCh      chan struct{}
	m             sync.Mutex
	finishOnce    sync.Once
}

// Add progress bars.
func (p *Pool) Add(pbs ...*ProgressBar) {
	p.m.Lock()
	defer p.m.Unlock()
	for _, bar := range pbs {
		bar.ManualUpdate = true
		bar.NotPrint = true
		bar.Start()
		p.bars = append(p.bars, bar)
	}
}

func (p *Pool) Start() (err error) {
	p.RefreshRate = DefaultRefreshRate
	p.shutdownCh, err = lockEcho()
	if err != nil {
		return
	}
	p.workerCh = make(chan struct{})
	go p.writer()
	return
}

func (p *Pool) writer() {
	var first = true
	defer func() {
		if first == false {
			p.print(false)
		} else {
			p.print(true)
			p.print(false)
		}
		close(p.workerCh)
	}()

	for {
		select {
		case <-time.After(p.RefreshRate):
			if p.print(first) {
				p.print(false)
				return
			}
			first = false
		case <-p.shutdownCh:
			return
		}
	}
}

// Restore terminal state and close pool
func (p *Pool) Stop() error {
	p.finishOnce.Do(func() {
		close(p.shutdownCh)
	})

	// Wait for the worker to complete
	select {
	case <-p.workerCh:
	}

	return unlockEcho()
}
<|endoftext|>


<commit_before>package worker

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/context"
)

const (
	DefaultWorkersCount = 10
)

var (
	DefaultTTR time.Duration = 60 * time.Second
)

type response struct {
	Msg *Message
	Err error
}

// Pool represents a pool of workers connected to a queue.
type Pool struct {
	count int   // workers count
	queue Queue // input queue

	middleware middleware
	handlers   []Handler
	mux        map[string]Factory
	logger     *log.Logger
}

// SetQueue assigns a custom queue to worker pool.
func SetQueue(q Queue) func(*Pool) {
	return func(p *Pool) {
		p.queue = q
	}
}

// SetWorkers configures the pool concurrency.
func SetWorkers(n int) func(*Pool) {
	return func(p *Pool) {
		p.count = n
	}
}

// NewPool returns a new Pool instance.
func NewPool(opts ...func(*Pool)) *Pool {
	pool := &Pool{
		count:    DefaultWorkersCount,
		queue:    NewMemoryQueue(),
		mux:      map[string]Factory{},
		logger:   log.New(os.Stdout, "[worker] ", 0),
		handlers: CommonStack(),
	}

	// Apply options.
	for _, opt := range opts {
		opt(pool)
	}

	// Init middleware stack.
	pool.middleware = pool.build(pool.handlers)

	return pool
}

// Add registers a new job factory.
func (p *Pool) Add(f Factory) error {
	typ, err := StructType(f)
	if err != nil {
		return err
	}

	if _, ok := p.mux[typ]; ok {
		return NewErrorFmt("factory %q exists already", typ)
	}
	p.mux[typ] = f
	return nil
}

// Use appends a new middleware to current stack.
func (p *Pool) Use(h Handler) {
	p.handlers = append(p.handlers, h)
	p.middleware = p.build(p.handlers)
}

// build iterates over the handlers, it returns a
// list of middlewares with each item linked to
// the next one.
func (p *Pool) build(hs []Handler) middleware {
	if len(hs) == 0 {
		return p.last()
	}
	next := p.build(hs[1:])
	return middleware{hs[0], &next}
}

// last builds and returns the last middleware, which will
// execute the job without calling next middleware.
func (p *Pool) last() middleware {
	return middleware{
		HandlerFunc(func(sw StatusWriter, fact string, args *Args, next JobRunner) {
			err := p.execute(fact, args)
			sw.Set(err)
		}),
		&middleware{},
	}
}

// Run starts processing jobs from the queue.
func (p *Pool) Run(ctx context.Context) error {
	var wg sync.WaitGroup

	// Handle unix signals.
	sig := p.trap()

	ctx, cancel := context.WithCancel(ctx)

	// Fan-out channel.
	c := make(chan *Message)

	// Start workers.
	wg.Add(p.count + 1)
	for i := 0; i < p.count; i++ {
		go func() {
			defer wg.Done()
			p.worker(ctx, c)
		}()
	}

	// Start the master.
	go func() {
		defer wg.Done()
		p.master(ctx, c)
	}()

	var err error
	select {
	case <-ctx.Done():
		err = ctx.Err()
	case <-sig:
		cancel()
	}

	p.logger.Println("Stopping workers ...")
	close(c)
	wg.Wait()
	p.logger.Println("Shutdown completed!")
	return err
}

func (p *Pool) master(ctx context.Context, c chan<- *Message) {
	var r *response
	for {
		select {
		case <-ctx.Done():
			return
		case r = <-p.get():
		}

		if r.Err != nil {
			if wkerr, ok := r.Err.(*Error); ok && wkerr.Temporary() {
				continue
			}
			return
		}

		select {
		case <-ctx.Done():
			return
		case c <- r.Msg:
		}
	}
}

// worker executes jobs from the in channel in a separate goroutine.
func (p *Pool) worker(ctx context.Context, in <-chan *Message) {
	for msg := range in {
		status := NewStatusWriter()
		done := make(chan struct{}, 1)

		// Start the job in a separate goroutine.
		go func() {
			p.Exec(status, msg.Type(), msg.Args())
			done <- struct{}{}
		}()

		// Wait job completion.
		select {
		case <-ctx.Done():
			return
		case <-done:
			if status.OK() {
				if err := p.queue.Ack(msg); err != nil {
					p.logger.Println("delete:", msg, err)
				}
			}
		}
	}
}

// Exec runs the job passing it through the middleware stack.
func (p *Pool) Exec(sw StatusWriter, fact string, args *Args) {
	p.middleware.Exec(sw, fact, args)
}

// execute runs the job without passing it through the
// middleware stack.
func (p *Pool) execute(fact string, args *Args) error {
	f, ok := p.mux[fact]
	if !ok {
		return NewErrorFmt("bad type: %v", fact)
	}

	j, err := f.Make(args)
	if err != nil {
		return NewErrorFmt("make: %v", err)
	}

	if err := j.Run(); err != nil {
		return NewErrorFmt("run: %v", err)
	}

	return nil
}

// trap traps OS signals to allow clean exit.
func (p *Pool) trap() <-chan struct{} {
	out := make(chan struct{}, 1)
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGUSR1, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		defer close(out)
		for s := range signals {
			switch s {
			case syscall.SIGINT, syscall.SIGUSR1, syscall.SIGTERM:
				p.logger.Println("Quit signal received ...")
				out <- struct{}{}
				return
			}
		}
	}()

	return out
}

// get service peeks a message from queue and
// returns it through a channel without blocking
// the caller.
func (p *Pool) get() <-chan *response {
	// Buffered channel to avoid goroutine leaking.
	out := make(chan *response, 1)

	// Start a gorouting to avoid blocking.
	go func() {
		defer close(out)

		msg, err := p.queue.Get()
		if err != nil {
			out <- &response{Err: err}
			return
		}

		out <- &response{Msg: msg}
	}()

	return out
}
<commit_msg>Update log message<commit_after>package worker

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/context"
)

const (
	DefaultWorkersCount = 10
)

var (
	DefaultTTR time.Duration = 60 * time.Second
)

type response struct {
	Msg *Message
	Err error
}

// Pool represents a pool of workers connected to a queue.
type Pool struct {
	count int   // workers count
	queue Queue // input queue

	middleware middleware
	handlers   []Handler
	mux        map[string]Factory
	logger     *log.Logger
}

// SetQueue assigns a custom queue to worker pool.
func SetQueue(q Queue) func(*Pool) {
	return func(p *Pool) {
		p.queue = q
	}
}

// SetWorkers configures the pool concurrency.
func SetWorkers(n int) func(*Pool) {
	return func(p *Pool) {
		p.count = n
	}
}

// NewPool returns a new Pool instance.
func NewPool(opts ...func(*Pool)) *Pool {
	pool := &Pool{
		count:    DefaultWorkersCount,
		queue:    NewMemoryQueue(),
		mux:      map[string]Factory{},
		logger:   log.New(os.Stdout, "[worker] ", 0),
		handlers: CommonStack(),
	}

	// Apply options.
	for _, opt := range opts {
		opt(pool)
	}

	// Init middleware stack.
	pool.middleware = pool.build(pool.handlers)

	return pool
}

// Add registers a new job factory.
func (p *Pool) Add(f Factory) error {
	typ, err := StructType(f)
	if err != nil {
		return err
	}

	if _, ok := p.mux[typ]; ok {
		return NewErrorFmt("factory %q exists already", typ)
	}
	p.mux[typ] = f
	return nil
}

// Use appends a new middleware to current stack.
func (p *Pool) Use(h Handler) {
	p.handlers = append(p.handlers, h)
	p.middleware = p.build(p.handlers)
}

// build iterates over the handlers, it returns a
// list of middlewares with each item linked to
// the next one.
func (p *Pool) build(hs []Handler) middleware {
	if len(hs) == 0 {
		return p.last()
	}
	next := p.build(hs[1:])
	return middleware{hs[0], &next}
}

// last builds and returns the last middleware, which will
// execute the job without calling next middleware.
func (p *Pool) last() middleware {
	return middleware{
		HandlerFunc(func(sw StatusWriter, fact string, args *Args, next JobRunner) {
			err := p.execute(fact, args)
			sw.Set(err)
		}),
		&middleware{},
	}
}

// Run starts processing jobs from the queue.
func (p *Pool) Run(ctx context.Context) error {
	var wg sync.WaitGroup

	// Handle unix signals.
	sig := p.trap()

	ctx, cancel := context.WithCancel(ctx)

	// Fan-out channel.
	c := make(chan *Message)

	// Start workers.
	wg.Add(p.count + 1)
	for i := 0; i < p.count; i++ {
		go func() {
			defer wg.Done()
			p.worker(ctx, c)
		}()
	}

	// Start the master.
	go func() {
		defer wg.Done()
		p.master(ctx, c)
	}()

	var err error
	select {
	case <-ctx.Done():
		err = ctx.Err()
	case <-sig:
		cancel()
	}

	p.logger.Println("Stopping workers ...")
	close(c)
	wg.Wait()
	p.logger.Println("Shutdown completed!")
	return err
}

func (p *Pool) master(ctx context.Context, c chan<- *Message) {
	var r *response
	for {
		select {
		case <-ctx.Done():
			return
		case r = <-p.get():
		}

		if r.Err != nil {
			if wkerr, ok := r.Err.(*Error); ok && wkerr.Temporary() {
				continue
			}
			return
		}

		select {
		case <-ctx.Done():
			return
		case c <- r.Msg:
		}
	}
}

// worker executes jobs from the in channel in a separate goroutine.
func (p *Pool) worker(ctx context.Context, in <-chan *Message) {
	for msg := range in {
		status := NewStatusWriter()
		done := make(chan struct{}, 1)

		// Start the job in a separate goroutine.
		go func() {
			p.Exec(status, msg.Type(), msg.Args())
			done <- struct{}{}
		}()

		// Wait job completion.
		select {
		case <-ctx.Done():
			return
		case <-done:
			if status.OK() {
				if err := p.queue.Ack(msg); err != nil {
					p.logger.Println("Ack failure:", msg, err)
				}
			}
		}
	}
}

// Exec runs the job passing it through the middleware stack.
func (p *Pool) Exec(sw StatusWriter, fact string, args *Args) {
	p.middleware.Exec(sw, fact, args)
}

// execute runs the job without passing it through the
// middleware stack.
func (p *Pool) execute(fact string, args *Args) error {
	f, ok := p.mux[fact]
	if !ok {
		return NewErrorFmt("bad type: %v", fact)
	}

	j, err := f.Make(args)
	if err != nil {
		return NewErrorFmt("make: %v", err)
	}

	if err := j.Run(); err != nil {
		return NewErrorFmt("run: %v", err)
	}

	return nil
}

// trap traps OS signals to allow clean exit.
func (p *Pool) trap() <-chan struct{} {
	out := make(chan struct{}, 1)
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGUSR1, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		defer close(out)
		for s := range signals {
			switch s {
			case syscall.SIGINT, syscall.SIGUSR1, syscall.SIGTERM:
				p.logger.Println("Quit signal received ...")
				out <- struct{}{}
				return
			}
		}
	}()

	return out
}

// get service peeks a message from queue and
// returns it through a channel without blocking
// the caller.
func (p *Pool) get() <-chan *response {
	// Buffered channel to avoid goroutine leaking.
	out := make(chan *response, 1)

	// Start a gorouting to avoid blocking.
	go func() {
		defer close(out)

		msg, err := p.queue.Get()
		if err != nil {
			out <- &response{Err: err}
			return
		}

		out <- &response{Msg: msg}
	}()

	return out
}
<|endoftext|>


<commit_before>package bytebufferpool

import (
	"sort"
	"sync"
	"sync/atomic"
)

const (
	minBitSize = 6 // CPU cache line size
	steps      = 20

	minSize = 1 << minBitSize
	maxSize = 1 << (minBitSize + steps - 1)

	calibrateCallsThreshold = 42000
	maxPercentile           = 0.95
)

// Pool represents byte buffer pool.
//
// Distinct pools may be used for distinct types of byte buffers.
// Properly determined byte buffer types with their own pools may help reducing
// memory waste.
type Pool struct {
	calls       [steps]uint64
	calibrating uint64

	defaultSize uint64
	maxSize     uint64

	pool sync.Pool
}

// Get returns new byte buffer with zero length.
//
// The byte buffer may be returned to the pool via Put after the use
// in order to minimize GC overhead.
func (p *Pool) Get() *ByteBuffer {
	v := p.pool.Get()
	if v != nil {
		return v.(*ByteBuffer)
	}
	return &ByteBuffer{
		B: make([]byte, 0, atomic.LoadUint64(&p.defaultSize)),
	}
}

// Put releases byte buffer obtained via Get to the pool.
//
// The bufer mustn't be accessed after returning to the pool.
func (p *Pool) Put(b *ByteBuffer) {
	idx := index(len(b.B))

	if atomic.AddUint64(&p.calls[idx], 1) > calibrateCallsThreshold {
		p.calibrate()
	}

	maxSize := int(atomic.LoadUint64(&p.maxSize))
	if maxSize == 0 || cap(b.B) <= maxSize {
		b.B = b.B[:0]
		p.pool.Put(b)
	}
}

func (p *Pool) calibrate() {
	if !atomic.CompareAndSwapUint64(&p.calibrating, 0, 1) {
		return
	}

	a := make(callSizes, 0, steps)
	var callsSum uint64
	for i := uint64(0); i < steps; i++ {
		calls := atomic.SwapUint64(&p.calls[i], 0)
		callsSum += calls
		a = append(a, callSize{
			calls: calls,
			size:  minSize << i,
		})
	}
	sort.Sort(a)

	defaultSize := a[0].size
	maxSize := defaultSize

	maxSum := uint64(float64(callsSum) * maxPercentile)
	callsSum = 0
	for i := 0; i < steps; i++ {
		if callsSum > maxSum {
			break
		}
		callsSum += a[i].calls
		size := a[i].size
		if size > maxSize {
			maxSize = size
		}
	}

	atomic.StoreUint64(&p.defaultSize, defaultSize)
	atomic.StoreUint64(&p.maxSize, maxSize)

	atomic.StoreUint64(&p.calibrating, 0)
}

type callSize struct {
	calls uint64
	size  uint64
}

type callSizes []callSize

func (ci callSizes) Len() int {
	return len(ci)
}

func (ci callSizes) Less(i, j int) bool {
	return ci[i].calls > ci[j].calls
}

func (ci callSizes) Swap(i, j int) {
	ci[i], ci[j] = ci[j], ci[i]
}

func index(n int) int {
	n--
	n >>= minBitSize
	idx := 0
	for n > 0 {
		n >>= 1
		idx++
	}
	if idx >= steps {
		idx = steps - 1
	}
	return idx
}
<commit_msg>Issue #1: clarified the comment regarding CPU cache line size<commit_after>package bytebufferpool

import (
	"sort"
	"sync"
	"sync/atomic"
)

const (
	minBitSize = 6 // 2**6=64 is a CPU cache line size
	steps      = 20

	minSize = 1 << minBitSize
	maxSize = 1 << (minBitSize + steps - 1)

	calibrateCallsThreshold = 42000
	maxPercentile           = 0.95
)

// Pool represents byte buffer pool.
//
// Distinct pools may be used for distinct types of byte buffers.
// Properly determined byte buffer types with their own pools may help reducing
// memory waste.
type Pool struct {
	calls       [steps]uint64
	calibrating uint64

	defaultSize uint64
	maxSize     uint64

	pool sync.Pool
}

// Get returns new byte buffer with zero length.
//
// The byte buffer may be returned to the pool via Put after the use
// in order to minimize GC overhead.
func (p *Pool) Get() *ByteBuffer {
	v := p.pool.Get()
	if v != nil {
		return v.(*ByteBuffer)
	}
	return &ByteBuffer{
		B: make([]byte, 0, atomic.LoadUint64(&p.defaultSize)),
	}
}

// Put releases byte buffer obtained via Get to the pool.
//
// The bufer mustn't be accessed after returning to the pool.
func (p *Pool) Put(b *ByteBuffer) {
	idx := index(len(b.B))

	if atomic.AddUint64(&p.calls[idx], 1) > calibrateCallsThreshold {
		p.calibrate()
	}

	maxSize := int(atomic.LoadUint64(&p.maxSize))
	if maxSize == 0 || cap(b.B) <= maxSize {
		b.B = b.B[:0]
		p.pool.Put(b)
	}
}

func (p *Pool) calibrate() {
	if !atomic.CompareAndSwapUint64(&p.calibrating, 0, 1) {
		return
	}

	a := make(callSizes, 0, steps)
	var callsSum uint64
	for i := uint64(0); i < steps; i++ {
		calls := atomic.SwapUint64(&p.calls[i], 0)
		callsSum += calls
		a = append(a, callSize{
			calls: calls,
			size:  minSize << i,
		})
	}
	sort.Sort(a)

	defaultSize := a[0].size
	maxSize := defaultSize

	maxSum := uint64(float64(callsSum) * maxPercentile)
	callsSum = 0
	for i := 0; i < steps; i++ {
		if callsSum > maxSum {
			break
		}
		callsSum += a[i].calls
		size := a[i].size
		if size > maxSize {
			maxSize = size
		}
	}

	atomic.StoreUint64(&p.defaultSize, defaultSize)
	atomic.StoreUint64(&p.maxSize, maxSize)

	atomic.StoreUint64(&p.calibrating, 0)
}

type callSize struct {
	calls uint64
	size  uint64
}

type callSizes []callSize

func (ci callSizes) Len() int {
	return len(ci)
}

func (ci callSizes) Less(i, j int) bool {
	return ci[i].calls > ci[j].calls
}

func (ci callSizes) Swap(i, j int) {
	ci[i], ci[j] = ci[j], ci[i]
}

func index(n int) int {
	n--
	n >>= minBitSize
	idx := 0
	for n > 0 {
		n >>= 1
		idx++
	}
	if idx >= steps {
		idx = steps - 1
	}
	return idx
}
<|endoftext|>


<commit_before>// Package pop3 implements the Post Office Protocol version 3.
package pop3

import (
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
)

var (
	EOF = errors.New("skip the all mail remaining")
)

// MessageInfo has Number, Size, and Uid fields,
// and used as a return value of ListAll and UidlAll.
// When used as the return value of the method ListAll,
// MessageInfo contain only the Number and Size values.
// When used as the return value of the method UidlAll,
// MessageInfo contain only the Number and Uid values.
type MessageInfo struct {
	Number int
	Size   uint64
	Uid    string
}

// A Client represents a client connection to an POP server.
type Client struct {
	// Text is the pop3.Conn used by the Client.
	Text *Conn
	// keep a reference to the connection so it can be used to create a TLS
	// connection later
	conn net.Conn
}

// Dial returns a new Client connected to an POP server at addr.
// The addr must include a port number.
func Dial(addr string) (*Client, error) {
	conn, err := net.Dial("tcp", addr)

	if err != nil {
		return nil, err
	}

	return NewClient(conn)
}

// NewClient returns a new Client using an existing connection.
func NewClient(conn net.Conn) (*Client, error) {
	text := NewConn(conn)

	_, err := text.ReadResponse()

	if err != nil {
		return nil, err
	}

	return &Client{Text: text, conn: conn}, nil
}

// User issues a USER command to the server using the provided user name.
func (c *Client) User(user string) error {
	return c.cmdSimple("USER %s", user)
}

// Pass issues a PASS command to the server using the provided password.
func (c *Client) Pass(pass string) error {
	return c.cmdSimple("PASS %s", pass)
}

// Stat issues a STAT command to the server
// and returns mail count and total size.
func (c *Client) Stat() (int, uint64, error) {
	return c.cmdStatOrList("STAT", "STAT")
}

// Retr issues a RETR command to the server using the provided mail number
// and returns mail data.
func (c *Client) Retr(number int) (string, error) {
	var err error

	err = c.Text.WriteLine("RETR %d", number)

	if err != nil {
		return "", err
	}

	_, err = c.Text.ReadResponse()

	if err != nil {
		return "", err
	}

	return c.Text.ReadToPeriod()
}

// List issues a LIST command to the server using the provided mail number
// and returns mail number and size.
func (c *Client) List(number int) (int, uint64, error) {
	return c.cmdStatOrList("LIST", "LIST %d", number)
}

// List issues a LIST command to the server
// and returns array of MessageInfo.
func (c *Client) ListAll() ([]MessageInfo, error) {
	list := make([]MessageInfo, 0)

	err := c.cmdReadLines("LIST", func(line string) error {
		number, size, err := c.convertNumberAndSize(line)

		if err != nil {
			return err
		}

		list = append(list, MessageInfo{Number: number, Size: size})

		return nil
	})

	if err != nil {
		return nil, err
	}

	return list, nil
}

// Uidl issues a UIDL command to the server using the provided mail number
// and returns mail number and unique id.
func (c *Client) Uidl(number int) (int, string, error) {
	var err error

	err = c.Text.WriteLine("UIDL %d", number)

	if err != nil {
		return 0, "", err
	}

	var msg string

	msg, err = c.Text.ReadResponse()

	if err != nil {
		return 0, "", err
	}

	var val int
	var uid string

	val, uid, err = c.convertNumberAndUid(msg)

	if err != nil {
		return 0, "", err
	}

	return val, uid, nil
}

// Uidl issues a UIDL command to the server
// and returns array of MessageInfo.
func (c *Client) UidlAll() ([]MessageInfo, error) {
	list := make([]MessageInfo, 0)

	err := c.cmdReadLines("UIDL", func(line string) error {
		number, uid, err := c.convertNumberAndUid(line)

		if err != nil {
			return err
		}

		list = append(list, MessageInfo{Number: number, Uid: uid})

		return nil
	})

	if err != nil {
		return nil, err
	}

	return list, nil
}

// Dele issues a DELE command to the server using the provided mail number.
func (c *Client) Dele(number int) error {
	return c.cmdSimple("DELE %d", number)
}

// Noop issues a NOOP command to the server.
func (c *Client) Noop() error {
	return c.cmdSimple("NOOP")
}

// Rset issues a RSET command to the server.
func (c *Client) Rset() error {
	return c.cmdSimple("RSET")
}

// Quit issues a QUIT command to the server.
func (c *Client) Quit() error {
	return c.cmdSimple("QUIT")
}

// ReceiveMail connects to the server at addr,
// and authenticates with user and pass,
// and calling receiveFn for each mail.
func ReceiveMail(addr, user, pass string, receiveFn ReceiveMailFunc) error {
	c, err := Dial(addr)

	if err != nil {
		return err
	}

	defer func() {
		c.Quit()
		c.Close()
	}()

	if err = c.User(user); err != nil {
		return err
	}

	if err = c.Pass(pass); err != nil {
		return err
	}

	var mis []MessageInfo

	if mis, err = c.UidlAll(); err != nil {
		return err
	}

	for _, mi := range mis {
		var data string

		data, err = c.Retr(mi.Number)

		del, err := receiveFn(mi.Number, mi.Uid, data, err)

		if err != nil && err != EOF {
			return err
		}

		if del {
			if err = c.Dele(mi.Number); err != nil {
				return err
			}
		}

		if err == EOF {
			break
		}
	}

	return nil
}

// ReceiveMailFunc is the type of the function called for each mail.
// Its arguments are mail's number, uid, data, and mail receiving error.
// if this function returns false value, the mail will be deleted,
// if its returns EOF, skip the all mail of remaining.
// (after deleting mail, if necessary)
type ReceiveMailFunc func(number int, uid, data string, err error) (bool, error)

func (c *Client) cmdSimple(format string, args ...interface{}) error {
	var err error

	err = c.Text.WriteLine(format, args...)

	if err != nil {
		return err
	}

	_, err = c.Text.ReadResponse()

	if err != nil {
		return err
	}

	return nil
}

func (c *Client) cmdStatOrList(name, format string, args ...interface{}) (int, uint64, error) {
	var err error

	err = c.Text.WriteLine(format, args...)

	if err != nil {
		return 0, 0, err
	}

	var msg string

	msg, err = c.Text.ReadResponse()

	if err != nil {
		return 0, 0, err
	}

	s := strings.Split(msg, " ")

	if len(s) < 2 {
		return 0, 0, ResponseError(fmt.Sprintf("invalid response format: %s", msg))
	}

	var val int
	var size uint64

	val, size, err = c.convertNumberAndSize(msg)

	if err != nil {
		return 0, 0, err
	}

	return val, size, nil
}

func (c *Client) cmdReadLines(cmnd string, lineFn lineFunc) error {
	var err error

	err = c.Text.WriteLine(cmnd)

	if err != nil {
		return err
	}

	_, err = c.Text.ReadResponse()

	if err != nil {
		return err
	}

	var lines []string

	lines, err = c.Text.ReadLines()

	if err != nil {
		return err
	}

	for _, line := range lines {
		err = lineFn(line)

		if err != nil {
			return err
		}
	}

	return nil
}

type lineFunc func(line string) error

func (c *Client) Close() error {
	return c.Text.Close()
}

func (c *Client) convertNumberAndSize(line string) (int, uint64, error) {
	var err error

	s := strings.Split(line, " ")

	if len(s) < 2 {
		return 0, 0, errors.New(fmt.Sprintf("the length of the array is less than 2: %s", line))
	}

	var val int
	var size uint64

	if val, err = strconv.Atoi(s[0]); err != nil {
		return 0, 0, errors.New(fmt.Sprintf("can not convert element[0] to int type: %s", line))
	}

	if size, err = strconv.ParseUint(s[1], 10, 64); err != nil {
		return 0, 0, errors.New(fmt.Sprintf("can not convert element[1] to uint64 type: %s", line))
	}

	return val, size, nil
}

func (c *Client) convertNumberAndUid(line string) (int, string, error) {
	var err error

	s := strings.Split(line, " ")

	if len(s) < 2 {
		return 0, "", errors.New(fmt.Sprintf("the length of the array is less than 2: %s", line))
	}

	var val int

	if val, err = strconv.Atoi(s[0]); err != nil {
		return 0, "", errors.New(fmt.Sprintf("can not convert element[0] to int type: %s", line))
	}

	return val, s[1], nil
}
<commit_msg>fix bug of ReceiveMail<commit_after>// Package pop3 implements the Post Office Protocol version 3.
package pop3

import (
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
)

var (
	EOF = errors.New("skip the all mail remaining")
)

// MessageInfo has Number, Size, and Uid fields,
// and used as a return value of ListAll and UidlAll.
// When used as the return value of the method ListAll,
// MessageInfo contain only the Number and Size values.
// When used as the return value of the method UidlAll,
// MessageInfo contain only the Number and Uid values.
type MessageInfo struct {
	Number int
	Size   uint64
	Uid    string
}

// A Client represents a client connection to an POP server.
type Client struct {
	// Text is the pop3.Conn used by the Client.
	Text *Conn
	// keep a reference to the connection so it can be used to create a TLS
	// connection later
	conn net.Conn
}

// Dial returns a new Client connected to an POP server at addr.
// The addr must include a port number.
func Dial(addr string) (*Client, error) {
	conn, err := net.Dial("tcp", addr)

	if err != nil {
		return nil, err
	}

	return NewClient(conn)
}

// NewClient returns a new Client using an existing connection.
func NewClient(conn net.Conn) (*Client, error) {
	text := NewConn(conn)

	_, err := text.ReadResponse()

	if err != nil {
		return nil, err
	}

	return &Client{Text: text, conn: conn}, nil
}

// User issues a USER command to the server using the provided user name.
func (c *Client) User(user string) error {
	return c.cmdSimple("USER %s", user)
}

// Pass issues a PASS command to the server using the provided password.
func (c *Client) Pass(pass string) error {
	return c.cmdSimple("PASS %s", pass)
}

// Stat issues a STAT command to the server
// and returns mail count and total size.
func (c *Client) Stat() (int, uint64, error) {
	return c.cmdStatOrList("STAT", "STAT")
}

// Retr issues a RETR command to the server using the provided mail number
// and returns mail data.
func (c *Client) Retr(number int) (string, error) {
	var err error

	err = c.Text.WriteLine("RETR %d", number)

	if err != nil {
		return "", err
	}

	_, err = c.Text.ReadResponse()

	if err != nil {
		return "", err
	}

	return c.Text.ReadToPeriod()
}

// List issues a LIST command to the server using the provided mail number
// and returns mail number and size.
func (c *Client) List(number int) (int, uint64, error) {
	return c.cmdStatOrList("LIST", "LIST %d", number)
}

// List issues a LIST command to the server
// and returns array of MessageInfo.
func (c *Client) ListAll() ([]MessageInfo, error) {
	list := make([]MessageInfo, 0)

	err := c.cmdReadLines("LIST", func(line string) error {
		number, size, err := c.convertNumberAndSize(line)

		if err != nil {
			return err
		}

		list = append(list, MessageInfo{Number: number, Size: size})

		return nil
	})

	if err != nil {
		return nil, err
	}

	return list, nil
}

// Uidl issues a UIDL command to the server using the provided mail number
// and returns mail number and unique id.
func (c *Client) Uidl(number int) (int, string, error) {
	var err error

	err = c.Text.WriteLine("UIDL %d", number)

	if err != nil {
		return 0, "", err
	}

	var msg string

	msg, err = c.Text.ReadResponse()

	if err != nil {
		return 0, "", err
	}

	var val int
	var uid string

	val, uid, err = c.convertNumberAndUid(msg)

	if err != nil {
		return 0, "", err
	}

	return val, uid, nil
}

// Uidl issues a UIDL command to the server
// and returns array of MessageInfo.
func (c *Client) UidlAll() ([]MessageInfo, error) {
	list := make([]MessageInfo, 0)

	err := c.cmdReadLines("UIDL", func(line string) error {
		number, uid, err := c.convertNumberAndUid(line)

		if err != nil {
			return err
		}

		list = append(list, MessageInfo{Number: number, Uid: uid})

		return nil
	})

	if err != nil {
		return nil, err
	}

	return list, nil
}

// Dele issues a DELE command to the server using the provided mail number.
func (c *Client) Dele(number int) error {
	return c.cmdSimple("DELE %d", number)
}

// Noop issues a NOOP command to the server.
func (c *Client) Noop() error {
	return c.cmdSimple("NOOP")
}

// Rset issues a RSET command to the server.
func (c *Client) Rset() error {
	return c.cmdSimple("RSET")
}

// Quit issues a QUIT command to the server.
func (c *Client) Quit() error {
	return c.cmdSimple("QUIT")
}

// ReceiveMail connects to the server at addr,
// and authenticates with user and pass,
// and calling receiveFn for each mail.
func ReceiveMail(addr, user, pass string, receiveFn ReceiveMailFunc) error {
	c, err := Dial(addr)

	if err != nil {
		return err
	}

	defer func() {
		if err != nil && err != EOF {
			c.Rset()
		}

		c.Quit()
		c.Close()
	}()

	if err = c.User(user); err != nil {
		return err
	}

	if err = c.Pass(pass); err != nil {
		return err
	}

	var mis []MessageInfo

	if mis, err = c.UidlAll(); err != nil {
		return err
	}

	for _, mi := range mis {
		var data string

		data, err = c.Retr(mi.Number)

		del, err := receiveFn(mi.Number, mi.Uid, data, err)

		if err != nil && err != EOF {
			return err
		}

		if del {
			if err = c.Dele(mi.Number); err != nil {
				return err
			}
		}

		if err == EOF {
			break
		}
	}

	return nil
}

// ReceiveMailFunc is the type of the function called for each mail.
// Its arguments are mail's number, uid, data, and mail receiving error.
// if this function returns false value, the mail will be deleted,
// if its returns EOF, skip the all mail of remaining.
// (after deleting mail, if necessary)
type ReceiveMailFunc func(number int, uid, data string, err error) (bool, error)

func (c *Client) cmdSimple(format string, args ...interface{}) error {
	var err error

	err = c.Text.WriteLine(format, args...)

	if err != nil {
		return err
	}

	_, err = c.Text.ReadResponse()

	if err != nil {
		return err
	}

	return nil
}

func (c *Client) cmdStatOrList(name, format string, args ...interface{}) (int, uint64, error) {
	var err error

	err = c.Text.WriteLine(format, args...)

	if err != nil {
		return 0, 0, err
	}

	var msg string

	msg, err = c.Text.ReadResponse()

	if err != nil {
		return 0, 0, err
	}

	s := strings.Split(msg, " ")

	if len(s) < 2 {
		return 0, 0, ResponseError(fmt.Sprintf("invalid response format: %s", msg))
	}

	var val int
	var size uint64

	val, size, err = c.convertNumberAndSize(msg)

	if err != nil {
		return 0, 0, err
	}

	return val, size, nil
}

func (c *Client) cmdReadLines(cmnd string, lineFn lineFunc) error {
	var err error

	err = c.Text.WriteLine(cmnd)

	if err != nil {
		return err
	}

	_, err = c.Text.ReadResponse()

	if err != nil {
		return err
	}

	var lines []string

	lines, err = c.Text.ReadLines()

	if err != nil {
		return err
	}

	for _, line := range lines {
		err = lineFn(line)

		if err != nil {
			return err
		}
	}

	return nil
}

type lineFunc func(line string) error

func (c *Client) Close() error {
	return c.Text.Close()
}

func (c *Client) convertNumberAndSize(line string) (int, uint64, error) {
	var err error

	s := strings.Split(line, " ")

	if len(s) < 2 {
		return 0, 0, errors.New(fmt.Sprintf("the length of the array is less than 2: %s", line))
	}

	var val int
	var size uint64

	if val, err = strconv.Atoi(s[0]); err != nil {
		return 0, 0, errors.New(fmt.Sprintf("can not convert element[0] to int type: %s", line))
	}

	if size, err = strconv.ParseUint(s[1], 10, 64); err != nil {
		return 0, 0, errors.New(fmt.Sprintf("can not convert element[1] to uint64 type: %s", line))
	}

	return val, size, nil
}

func (c *Client) convertNumberAndUid(line string) (int, string, error) {
	var err error

	s := strings.Split(line, " ")

	if len(s) < 2 {
		return 0, "", errors.New(fmt.Sprintf("the length of the array is less than 2: %s", line))
	}

	var val int

	if val, err = strconv.Atoi(s[0]); err != nil {
		return 0, "", errors.New(fmt.Sprintf("can not convert element[0] to int type: %s", line))
	}

	return val, s[1], nil
}
<|endoftext|>


<commit_before>// xlattice_go/prng.go
package rnglib

import "os"
import "math/rand"
import "strings"

type PRNG struct {
	rng *rand.Rand
}

//func NewMTSource(seed int64) rand.Source {
//	var mt64 MT64
//	mt64.Seed(seed)
//	return &mt64
//}
//
//func NewSimpleRNG(seed int64) *PRNG {
//	s := new(PRNG) // allocates
//	src := NewMTSource(seed)
//	s.rng = rand.New(src)
//	s.Seed(seed)
//	return s
//} // GEEP

func (s *PRNG) Seed(seed int64) {
	s.rng.Seed(seed)
}

// expose the rand.Random interface /////////////////////////////////
func (s *PRNG) Int63() int64         { return s.rng.Int63() }
func (s *PRNG) Uint32() uint32       { return s.rng.Uint32() }
func (s *PRNG) Int31() int32         { return s.rng.Int31() }
func (s *PRNG) Int() int             { return s.rng.Int() }
func (s *PRNG) Int63n(n int64) int64 { return s.rng.Int63n(n) }
func (s *PRNG) Int31n(n int32) int32 { return s.rng.Int31n(n) }
func (s *PRNG) Intn(n int) int       { return s.rng.Intn(n) }
func (s *PRNG) Float64() float64     { return s.rng.Float64() }
func (s *PRNG) Float32() float32     { return s.rng.Float32() }
func (s *PRNG) Perm(n int) []int     { return s.rng.Perm(n) }

// PRNG functions //////////////////////////////////////////////
func (s *PRNG) NextBoolean() bool { return s.rng.Int63()&1 == 0 }
func (s *PRNG) NextByte() byte    { return byte(s.rng.Int63n(256)) }

// miraculously inefficient
func (s *PRNG) NextBytes(buffer *[]byte) {
	for n := range *buffer {
		(*buffer)[n] = s.NextByte()
	}
}
func (s *PRNG) NextInt32(n uint32) uint32 { return uint32(float32(n) * s.rng.Float32()) }
func (s *PRNG) NextInt64(n uint64) uint64 { return uint64(float64(n) * s.rng.Float64()) }
func (s *PRNG) NextFloat32() float32      { return s.rng.Float32() }
func (s *PRNG) NextFloat64() float64      { return s.rng.Float64() }

// These produce strings which are acceptable POSIX file names.
// and also advance a cursor by a multiple of 64 bits.  All strings
// are at least one byte and less than maxLen bytes n length.  We
// arbitrarily limit file names to less than 256 characters.

func (s *PRNG) _nextFileName(nameLen int) string {
	/* always returns at least one character */
	maxStarterNdx := uint32(len(_FILE_NAME_STARTERS))
	ndx := s.NextInt32(maxStarterNdx)
	var chars []string
	chars = append(chars, _FILE_NAME_STARTERS[ndx])
	maxCharNdx := uint32(len(_FILE_NAME_CHARS))
	for n := 0; n < nameLen; n++ {
		ndx := s.NextInt32(maxCharNdx)
		chars = append(chars, _FILE_NAME_CHARS[ndx])
	}
	return strings.Join(chars, "")
}

func (s *PRNG) NextFileName(maxLen int) string {
	_maxLen := uint32(maxLen)
	if _maxLen < 2 {
		_maxLen = 2 // this is a ceiling which cannot be reached
	}
	if _maxLen > 256 {
		_maxLen = 256
	}
	var name string
	nameLen := uint32(0)
	for nameLen == 0 {
		nameLen = s.NextInt32(_maxLen) // so len < 256
	}
	for {
		name = s._nextFileName(int(nameLen))
		if (len(name) > 0) && !strings.Contains(name, "..") {
			break
		}
	}
	return name
}

// These are operations on the file system.  Directory depth is at least 1
// and no more than 'depth'.  Likewise for width, the number of
// files in a directory, where a file is either a data file or a subdirectory.
// The number of bytes in a file is at least minLen and less than maxLen.
// Subdirectory names may be random

func (s *PRNG) NextDataFile(dirName string, maxLen int, minLen int) (int, string) {
	// silently convert parameters to reasonable values
	if minLen < 0 {
		minLen = 0
	}
	if maxLen < minLen+1 {
		maxLen = minLen + 1
	}

	// create the data directory if it does not exist
	dirExists, err := PathExists(dirName)
	if err != nil {
		panic(err)
	}
	if !dirExists {
		os.MkdirAll(dirName, 0755)
	}

	// loop until the file does not exist
	pathToFile := dirName + "/" + s.NextFileName(16)
	pathExists, err := PathExists(pathToFile)
	if err != nil {
		panic(err)
	}
	for pathExists {
		pathToFile := dirName + "/" + s.NextFileName(16)
		pathExists, err = PathExists(pathToFile)
		if err != nil {
			panic(err)
		}
	}
	count := minLen + int(s.NextFloat32()*float32((maxLen-minLen)))
	data := make([]byte, count)
	s.NextBytes(&data) // fill with random bytes
	// XXX may cause panic
	fo, err := os.Create(pathToFile)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}() // XXX wakaranai
	// XXX this should be chunked
	// XXX data should be slice
	if _, err := fo.Write(data); err != nil {
		panic(err)
	}
	// XXX respec to also return err
	return count, pathToFile
}

// NextDataDir creates a directory tree populated with data files.
//
// BUGS
// * on at least one occasion with width = 4 only 3 files/directories
//   were created at the top level (2 were subdirs)
// DEFICIENCIES:
// * no control over percentage of directories
// * no guarantee that depth will be reached

func (s *PRNG) NextDataDir(pathToDir string, depth int, width int,
	maxLen int, minLen int) {
	// number of directory levels; 1 means no subdirectories
	if depth < 1 {
		depth = 1
	}
	// number of members (files, subdirectories) at each level
	if width < 1 {
		width = 1
	}
	// XXX may panic
	pathExists, err := PathExists(pathToDir)
	if err != nil {
		panic(err)
	}
	if !pathExists {
		os.MkdirAll(pathToDir, 0755)
	}
	subdirSoFar := 0
	for i := 0; i < width; i++ {
		if depth > 1 {
			if (s.NextFloat32() > 0.25) &&
				((i < width-1) || (subdirSoFar > 0)) {
				// 25% are subdirs
				// data file i
				// SPECIFICATION ERROR: file name may not be unique
				// count, pathToFile
				s.NextDataFile(pathToDir,
					maxLen, minLen)
			} else {
				// directory
				subdirSoFar += 1
				// create unique name
				fileName := s.NextFileName(16)
				pathToSubdir := pathToDir + "/" + fileName
				s.NextDataDir(pathToSubdir, depth-1, width,
					maxLen, minLen)
			}
		} else {
			// data file
			// XXX SPECIFICATION ERROR: file name may not be unique
			// count, pathToFile
			s.NextDataFile(pathToDir, maxLen, minLen)
		}
	} // end for
}
<commit_msg>working sha3_test.go<commit_after>// xlattice_go/prng.go
package rnglib

import "os"
import "math/rand"
import "strings"

type PRNG struct {
	rng *rand.Rand
}

//func NewMTSource(seed int64) rand.Source {
//	var mt64 MT64
//	mt64.Seed(seed)
//	return &mt64
//}
//
//func NewSimpleRNG(seed int64) *PRNG {
//	s := new(PRNG) // allocates
//	src := NewMTSource(seed)
//	s.rng = rand.New(src)
//	s.Seed(seed)
//	return s
//} // GEEP

func (s *PRNG) Seed(seed int64) {
	s.rng.Seed(seed)
}

// expose the rand.Random interface /////////////////////////////////
func (s *PRNG) Int63() int64         { return s.rng.Int63() }
func (s *PRNG) Uint32() uint32       { return s.rng.Uint32() }
func (s *PRNG) Int31() int32         { return s.rng.Int31() }
func (s *PRNG) Int() int             { return s.rng.Int() }
func (s *PRNG) Int63n(n int64) int64 { return s.rng.Int63n(n) }
func (s *PRNG) Int31n(n int32) int32 { return s.rng.Int31n(n) }
func (s *PRNG) Intn(n int) int       { return s.rng.Intn(n) }
func (s *PRNG) Float64() float64     { return s.rng.Float64() }
func (s *PRNG) Float32() float32     { return s.rng.Float32() }
func (s *PRNG) Perm(n int) []int     { return s.rng.Perm(n) }

// PRNG functions //////////////////////////////////////////////
func (s *PRNG) NextBoolean() bool { return s.rng.Int63()&1 == 0 }
func (s *PRNG) NextByte() byte    { return byte(s.rng.Int63n(256)) }

// miraculously inefficient
func (s *PRNG) NextBytes(buffer *[]byte) {
	for n := range *buffer {
		(*buffer)[n] = s.NextByte()
	}
}
func (s *PRNG) NextInt32(n uint32) uint32 { return uint32(float32(n) * s.rng.Float32()) }
func (s *PRNG) NextInt64(n uint64) uint64 { return uint64(float64(n) * s.rng.Float64()) }
func (s *PRNG) NextFloat32() float32      { return s.rng.Float32() }
func (s *PRNG) NextFloat64() float64      { return s.rng.Float64() }

// These produce strings which are acceptable POSIX file names.
// and also advance a cursor by a multiple of 64 bits.  All strings
// are at least one byte and less than maxLen bytes n length.  We
// arbitrarily limit file names to less than 256 characters.

func (s *PRNG) _nextFileName(nameLen int) string {
	/* always returns at least one character */
	maxStarterNdx := uint32(len(_FILE_NAME_STARTERS))
	ndx := s.NextInt32(maxStarterNdx)
	var chars []string
	chars = append(chars, _FILE_NAME_STARTERS[ndx])
	maxCharNdx := uint32(len(_FILE_NAME_CHARS))
	for n := 0; n < nameLen; n++ {
		ndx := s.NextInt32(maxCharNdx)
		chars = append(chars, _FILE_NAME_CHARS[ndx])
	}
	return strings.Join(chars, "")
}

func (s *PRNG) NextFileName(maxLen int) string {
	_maxLen := uint32(maxLen)
	if _maxLen < 2 {
		_maxLen = 2 // this is a ceiling which cannot be reached
	}
	if _maxLen > 256 {
		_maxLen = 256
	}
	var name string
	nameLen := uint32(0)
	for nameLen == 0 {
		nameLen = s.NextInt32(_maxLen) // so len < 256
	}
	for {
		name = s._nextFileName(int(nameLen))
		if (len(name) > 0) && !strings.Contains(name, "..") {
			break
		}
	}
	return name
}

// These are operations on the file system.  Directory depth is at least 1
// and no more than 'depth'.  Likewise for width, the number of
// files in a directory, where a file is either a data file or a subdirectory.
// The number of bytes in a file is at least minLen and less than maxLen.
// Subdirectory names may be random
//
// XXX CHANGED TO RETURN AN INT64 LENGTH

func (s *PRNG) NextDataFile(dirName string, maxLen int, minLen int) (int64, string) {
	// silently convert parameters to reasonable values
	if minLen < 0 {
		minLen = 0
	}
	if maxLen < minLen+1 {
		maxLen = minLen + 1
	}

	// create the data directory if it does not exist
	dirExists, err := PathExists(dirName)
	if err != nil {
		panic(err)
	}
	if !dirExists {
		os.MkdirAll(dirName, 0755)
	}

	// loop until the file does not exist
	pathToFile := dirName + "/" + s.NextFileName(16)
	pathExists, err := PathExists(pathToFile)
	if err != nil {
		panic(err)
	}
	for pathExists {
		pathToFile := dirName + "/" + s.NextFileName(16)
		pathExists, err = PathExists(pathToFile)
		if err != nil {
			panic(err)
		}
	}
	count := minLen + int(s.NextFloat32()*float32((maxLen-minLen)))
	data := make([]byte, count)
	s.NextBytes(&data) // fill with random bytes
	// XXX may cause panic
	fo, err := os.Create(pathToFile)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}() // XXX wakaranai
	// XXX this should be chunked
	// XXX data should be slice
	if _, err := fo.Write(data); err != nil {
		panic(err)
	}
	// XXX respec to also return err
	return int64(count), pathToFile
}

// NextDataDir creates a directory tree populated with data files.
//
// BUGS
// * on at least one occasion with width = 4 only 3 files/directories
//   were created at the top level (2 were subdirs)
// DEFICIENCIES:
// * no control over percentage of directories
// * no guarantee that depth will be reached

func (s *PRNG) NextDataDir(pathToDir string, depth int, width int,
	maxLen int, minLen int) {
	// number of directory levels; 1 means no subdirectories
	if depth < 1 {
		depth = 1
	}
	// number of members (files, subdirectories) at each level
	if width < 1 {
		width = 1
	}
	// XXX may panic
	pathExists, err := PathExists(pathToDir)
	if err != nil {
		panic(err)
	}
	if !pathExists {
		os.MkdirAll(pathToDir, 0755)
	}
	subdirSoFar := 0
	for i := 0; i < width; i++ {
		if depth > 1 {
			if (s.NextFloat32() > 0.25) &&
				((i < width-1) || (subdirSoFar > 0)) {
				// 25% are subdirs
				// data file i
				// SPECIFICATION ERROR: file name may not be unique
				// count, pathToFile
				s.NextDataFile(pathToDir,
					maxLen, minLen)
			} else {
				// directory
				subdirSoFar += 1
				// create unique name
				fileName := s.NextFileName(16)
				pathToSubdir := pathToDir + "/" + fileName
				s.NextDataDir(pathToSubdir, depth-1, width,
					maxLen, minLen)
			}
		} else {
			// data file
			// XXX SPECIFICATION ERROR: file name may not be unique
			// count, pathToFile
			s.NextDataFile(pathToDir, maxLen, minLen)
		}
	} // end for
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"syscall"

	"github.com/cenkalti/log"
	"github.com/cenkalti/rain/client"
	"github.com/cenkalti/rain/internal/logger"
	rainversion "github.com/cenkalti/rain/internal/version"
	"github.com/cenkalti/rain/resume/torrentresume"
	"github.com/cenkalti/rain/storage/filestorage"
	"github.com/cenkalti/rain/torrent"
	"github.com/mitchellh/go-homedir"
)

var (
	configPath = flag.String("config", "", "config path")
	dest       = flag.String("dest", ".", "where to download")
	port       = flag.Int("port", 0, "listen port")
	debug      = flag.Bool("debug", false, "enable debug log")
	version    = flag.Bool("version", false, "version")
	seed       = flag.Bool("seed", false, "continue seeding after dowload finishes")
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`")
)

func main() {
	flag.Parse()
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
		defer pprof.StopCPUProfile()
	}
	if *version {
		fmt.Println(rainversion.Version)
		return
	}
	args := flag.Args()
	if len(args) == 0 {
		_, _ = fmt.Fprintln(os.Stderr, "Give a torrent file as first argument!")
		os.Exit(1)
	}
	if *debug {
		logger.SetLogLevel(log.DEBUG)
	}
	cfg := client.NewConfig()
	if *configPath != "" {
		cp, err := homedir.Expand(*configPath)
		if err != nil {
			log.Fatal(err)
		}
		err = cfg.LoadFile(cp)
		if err != nil {
			log.Fatal(err)
		}
	}
	sto, err := filestorage.New(*dest)
	if err != nil {
		log.Fatal(err)
	}
	res, err := torrentresume.New("rain.resume")
	if err != nil {
		log.Fatal(err)
	}
	var t *torrent.Torrent
	if strings.HasPrefix(args[0], "magnet:") {
		t, err = torrent.DownloadMagnet(args[0], *port, sto, res)
	} else {
		f, err2 := os.Open(args[0])
		if err2 != nil {
			log.Fatal(err2)
		}
		t, err = torrent.DownloadTorrent(f, *port, sto, res)
		_ = f.Close()
	}
	if err != nil {
		log.Fatal(err)
	}

	sigC := make(chan os.Signal, 1)
	signal.Notify(sigC, syscall.SIGINT, syscall.SIGTERM)
LOOP:
	for {
		select {
		case <-sigC:
			break LOOP
		case <-t.NotifyComplete():
			if !*seed {
				break LOOP
			}
		case err = <-t.NotifyError():
			log.Error(err)
			os.Exit(1)
		}
	}
	t.Close()
}
<commit_msg>use new cli<commit_after>package main

import (
	"errors"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"syscall"

	"github.com/cenkalti/log"
	"github.com/cenkalti/rain/client"
	"github.com/cenkalti/rain/internal/logger"
	rainversion "github.com/cenkalti/rain/internal/version"
	"github.com/cenkalti/rain/resume/torrentresume"
	"github.com/cenkalti/rain/storage/filestorage"
	"github.com/cenkalti/rain/torrent"
	"github.com/mitchellh/go-homedir"
	"github.com/urfave/cli"
)

var (
	cfg = client.NewConfig()
	app = cli.NewApp()
)

func main() {
	app.Version = rainversion.Version
	app.Usage = "BitTorrent client"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Usage: "read config from `FILE`",
		},
		cli.StringFlag{
			Name:  "cpuprofile",
			Usage: "write cpu profile to `FILE`",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug log",
		},
	}
	app.Before = handleBeforeCommand
	app.After = handleAfterCommand
	app.Commands = []cli.Command{
		{
			Name:      "download",
			Usage:     "download torrent or magnet",
			ArgsUsage: "[torrent path or magnet link]",
			Action:    handleDownload,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "dest",
					Usage: "save files under `DIR`",
					Value: ".",
				},
				cli.IntFlag{
					Name:  "port",
					Usage: "peer listen port",
				},
				cli.BoolFlag{
					Name:  "seed",
					Usage: "continue seeding after download finishes",
				},
			},
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func handleBeforeCommand(c *cli.Context) error {
	cpuprofile := c.GlobalString("cpuprofile")
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
	}
	if c.GlobalBool("debug") {
		logger.SetLogLevel(log.DEBUG)
	}
	configPath := c.GlobalString("config")
	if configPath != "" {
		cp, err := homedir.Expand(configPath)
		if err != nil {
			log.Fatal(err)
		}
		err = cfg.LoadFile(cp)
		if err != nil {
			log.Fatal(err)
		}
	}
	return nil
}

func handleAfterCommand(c *cli.Context) error {
	if c.GlobalString("cpuprofile") != "" {
		defer pprof.StopCPUProfile()
	}
	return nil
}

func handleDownload(c *cli.Context) error {
	path := c.Args().Get(0)
	if path == "" {
		return errors.New("first argument must be a torrent file or magnet link")
	}
	sto, err := filestorage.New(c.String("dest"))
	if err != nil {
		log.Fatal(err)
	}
	res, err := torrentresume.New("rain.resume")
	if err != nil {
		log.Fatal(err)
	}
	var t *torrent.Torrent
	if strings.HasPrefix(path, "magnet:") {
		t, err = torrent.DownloadMagnet(path, c.Int("port"), sto, res)
	} else {
		f, err2 := os.Open(path) // nolint: gosec
		if err2 != nil {
			log.Fatal(err2)
		}
		t, err = torrent.DownloadTorrent(f, c.Int("port"), sto, res)
		_ = f.Close()
	}
	if err != nil {
		log.Fatal(err)
	}

	sigC := make(chan os.Signal, 1)
	signal.Notify(sigC, syscall.SIGINT, syscall.SIGTERM)
LOOP:
	for {
		select {
		case <-sigC:
			break LOOP
		case <-t.NotifyComplete():
			if !c.Bool("seed") {
				break LOOP
			}
		case err = <-t.NotifyError():
			log.Error(err)
			os.Exit(1)
		}
	}
	t.Close()
	return nil
}
<|endoftext|>