<commit_before>package workerpool

import (
	"github.com/gammazero/deque"
	"sync"
	"time"
)

const (
	// This value is the size of the queue that workers register their
	// availability to the dispatcher.  There may be hundreds of workers, but
	// only a small channel is needed to register some of the workers.
	readyQueueSize = 16

	// If worker pool receives no new work for this period of time, then stop
	// a worker goroutine.
	idleTimeoutSec = 5
)

// New creates and starts a pool of worker goroutines.
//
// The maxWorkers parameter specifies the maximum number of workers that will
// execute tasks concurrently.  After each timeout period, a worker goroutine
// is stopped until there are no remaining workers.
func New(maxWorkers int) *WorkerPool {
	// There must be at least one worker.
	if maxWorkers < 1 {
		maxWorkers = 1
	}

	pool := &WorkerPool{
		taskQueue:    make(chan func(), 1),
		maxWorkers:   maxWorkers,
		readyWorkers: make(chan chan func(), readyQueueSize),
		timeout:      time.Second * idleTimeoutSec,
		stoppedChan:  make(chan struct{}),
	}

	// Start the task dispatcher.
	go pool.dispatch()

	return pool
}

// WorkerPool is a collection of goroutines, where the number of concurrent
// goroutines processing requests does not exceed the specified maximum.
type WorkerPool struct {
	maxWorkers   int
	timeout      time.Duration
	taskQueue    chan func()
	readyWorkers chan chan func()
	stoppedChan  chan struct{}
	waitingQueue deque.Deque
	stopMutex    sync.Mutex
	stopped      bool
}

// Stop stops the worker pool and waits for only currently running tasks to
// complete.  Pending tasks that are not currently running are abandoned.
// Tasks must not be submitted to the worker pool after calling stop.
//
// Since creating the worker pool starts at least one goroutine, for the
// dispatcher, Stop() or StopWait() should be called when the worker pool is no
// longer needed.
func (p *WorkerPool) Stop() {
	p.stop(false)
}

// StopWait stops the worker pool and waits for all queued tasks tasks to
// complete.  No additional tasks may be submitted, but all pending tasks are
// executed by workers before this function returns.
func (p *WorkerPool) StopWait() {
	p.stop(true)
}

// Stopped returns true if this worker pool has been stopped.
func (p *WorkerPool) Stopped() bool {
	p.stopMutex.Lock()
	defer p.stopMutex.Unlock()
	return p.stopped
}

// Submit enqueues a function for a worker to execute.
//
// Any external values needed by the task function must be captured in a
// closure.  Any return values should be returned over a channel that is
// captured in the task function closure.
//
// Submit will not block regardless of the number of tasks submitted.  Each
// task is immediately given to an available worker or passed to a goroutine to
// be given to the next available worker.  If there are no available workers,
// the dispatcher adds a worker, until the maximum number of workers is
// running.
//
// After the maximum number of workers are running, and no workers are
// available, incoming tasks are put onto a queue and will be executed as
// workers become available.
//
// When no new tasks have been submitted for time period and a worker is
// available, the worker is shutdown.  As long as no new tasks arrive, one
// available worker is shutdown each time period until there are no more idle
// workers.  Since the time to start new goroutines is not significant, there
// is no need to retain idle workers.
func (p *WorkerPool) Submit(task func()) {
	if task != nil {
		p.taskQueue <- task
	}
}

// SubmitWait enqueues the given function and waits for it to be executed.
func (p *WorkerPool) SubmitWait(task func()) {
	if task == nil {
		return
	}
	doneChan := make(chan struct{})
	p.taskQueue <- func() {
		task()
		close(doneChan)
	}
	<-doneChan
}

// dispatch sends the next queued task to an available worker.
func (p *WorkerPool) dispatch() {
	defer close(p.stoppedChan)
	timeout := time.NewTimer(p.timeout)
	var (
		workerCount    int
		task           func()
		ok, wait       bool
		workerTaskChan chan func()
	)
	startReady := make(chan chan func())
Loop:
	for {
		// As long as tasks are in the waiting queue, remove and execute these
		// tasks as workers become available, and place new incoming tasks on
		// the queue.  Once the queue is empty, then go back to submitting
		// incoming tasks directly to available workers.
		if p.waitingQueue.Len() != 0 {
			select {
			case task, ok = <-p.taskQueue:
				if !ok {
					break Loop
				}
				if task == nil {
					wait = true
					break Loop
				}
				p.waitingQueue.PushBack(task)
			case workerTaskChan = <-p.readyWorkers:
				// A worker is ready, so give task to worker.
				workerTaskChan <- p.waitingQueue.PopFront().(func())
			}
			continue
		}
		timeout.Reset(p.timeout)
		select {
		case task, ok = <-p.taskQueue:
			if !ok || task == nil {
				break Loop
			}
			// Got a task to do.
			select {
			case workerTaskChan = <-p.readyWorkers:
				// A worker is ready, so give task to worker.
				workerTaskChan <- task
			default:
				// No workers ready.
				// Create a new worker, if not at max.
				if workerCount < p.maxWorkers {
					workerCount++
					go func(t func()) {
						startWorker(startReady, p.readyWorkers)
						// Submit the task when the new worker.
						taskChan := <-startReady
						taskChan <- t
					}(task)
				} else {
					// Enqueue task to be executed by next available worker.
					p.waitingQueue.PushBack(task)
				}
			}
		case <-timeout.C:
			// Timed out waiting for work to arrive.  Kill a ready worker.
			if workerCount > 0 {
				select {
				case workerTaskChan = <-p.readyWorkers:
					// A worker is ready, so kill.
					close(workerTaskChan)
					workerCount--
				default:
					// No work, but no ready workers.  All workers are busy.
				}
			}
		}
	}

	// If instructed to wait for all queued tasks, then remove from queue and
	// give to workers until queue is empty.
	if wait {
		for p.waitingQueue.Len() != 0 {
			workerTaskChan = <-p.readyWorkers
			// A worker is ready, so give task to worker.
			workerTaskChan <- p.waitingQueue.PopFront().(func())
		}
	}

	// Stop all remaining workers as they become ready.
	for workerCount > 0 {
		workerTaskChan = <-p.readyWorkers
		close(workerTaskChan)
		workerCount--
	}
}

// startWorker starts a goroutine that executes tasks given by the dispatcher.
//
// When a new worker starts, it registers its availability on the startReady
// channel.  This ensures that the goroutine associated with starting the
// worker gets to use the worker to execute its task.  Otherwise, the main
// dispatcher loop could steal the new worker and not know to start up another
// worker for the waiting goroutine.  The task would then have to wait for
// another existing worker to become available, even though capacity is
// available to start additional workers.
//
// A worker registers that is it available to do work by putting its task
// channel on the readyWorkers channel.  The dispatcher reads a worker's task
// channel from the readyWorkers channel, and writes a task to the worker over
// the worker's task channel.  To stop a worker, the dispatcher closes a
// worker's task channel, instead of writing a task to it.
func startWorker(startReady, readyWorkers chan chan func()) {
	go func() {
		taskChan := make(chan func())
		var task func()
		var ok bool
		// Register availability on starReady channel.
		startReady <- taskChan
		for {
			// Read task from dispatcher.
			task, ok = <-taskChan
			if !ok {
				// Dispatcher has told worker to stop.
				break
			}

			// Execute the task.
			task()

			// Register availability on readyWorkers channel.
			readyWorkers <- taskChan
		}
	}()
}

// stop tells the dispatcher to exit, and whether or not to complete queued
// tasks.
func (p *WorkerPool) stop(wait bool) {
	p.stopMutex.Lock()
	defer p.stopMutex.Unlock()
	if p.stopped {
		return
	}
	p.stopped = true
	if wait {
		p.taskQueue <- nil
	}
	// Close task queue and wait for currently running tasks to finish.
	close(p.taskQueue)
	<-p.stoppedChan
}
<commit_msg>An extremely simple method to return queue sizes<commit_after>package workerpool

import (
	"github.com/gammazero/deque"
	"sync"
	"time"
)

const (
	// This value is the size of the queue that workers register their
	// availability to the dispatcher.  There may be hundreds of workers, but
	// only a small channel is needed to register some of the workers.
	readyQueueSize = 16

	// If worker pool receives no new work for this period of time, then stop
	// a worker goroutine.
	idleTimeoutSec = 5
)

// New creates and starts a pool of worker goroutines.
//
// The maxWorkers parameter specifies the maximum number of workers that will
// execute tasks concurrently.  After each timeout period, a worker goroutine
// is stopped until there are no remaining workers.
func New(maxWorkers int) *WorkerPool {
	// There must be at least one worker.
	if maxWorkers < 1 {
		maxWorkers = 1
	}

	pool := &WorkerPool{
		taskQueue:    make(chan func(), 1),
		maxWorkers:   maxWorkers,
		readyWorkers: make(chan chan func(), readyQueueSize),
		timeout:      time.Second * idleTimeoutSec,
		stoppedChan:  make(chan struct{}),
	}

	// Start the task dispatcher.
	go pool.dispatch()

	return pool
}

// WorkerPool is a collection of goroutines, where the number of concurrent
// goroutines processing requests does not exceed the specified maximum.
type WorkerPool struct {
	maxWorkers   int
	timeout      time.Duration
	taskQueue    chan func()
	readyWorkers chan chan func()
	stoppedChan  chan struct{}
	waitingQueue deque.Deque
	stopMutex    sync.Mutex
	stopped      bool
}

// Stop stops the worker pool and waits for only currently running tasks to
// complete.  Pending tasks that are not currently running are abandoned.
// Tasks must not be submitted to the worker pool after calling stop.
//
// Since creating the worker pool starts at least one goroutine, for the
// dispatcher, Stop() or StopWait() should be called when the worker pool is no
// longer needed.
func (p *WorkerPool) Stop() {
	p.stop(false)
}

// StopWait stops the worker pool and waits for all queued tasks tasks to
// complete.  No additional tasks may be submitted, but all pending tasks are
// executed by workers before this function returns.
func (p *WorkerPool) StopWait() {
	p.stop(true)
}

// Stopped returns true if this worker pool has been stopped.
func (p *WorkerPool) Stopped() bool {
	p.stopMutex.Lock()
	defer p.stopMutex.Unlock()
	return p.stopped
}

// Submit enqueues a function for a worker to execute.
//
// Any external values needed by the task function must be captured in a
// closure.  Any return values should be returned over a channel that is
// captured in the task function closure.
//
// Submit will not block regardless of the number of tasks submitted.  Each
// task is immediately given to an available worker or passed to a goroutine to
// be given to the next available worker.  If there are no available workers,
// the dispatcher adds a worker, until the maximum number of workers is
// running.
//
// After the maximum number of workers are running, and no workers are
// available, incoming tasks are put onto a queue and will be executed as
// workers become available.
//
// When no new tasks have been submitted for time period and a worker is
// available, the worker is shutdown.  As long as no new tasks arrive, one
// available worker is shutdown each time period until there are no more idle
// workers.  Since the time to start new goroutines is not significant, there
// is no need to retain idle workers.
func (p *WorkerPool) Submit(task func()) {
	if task != nil {
		p.taskQueue <- task
	}
}

// SubmitWait enqueues the given function and waits for it to be executed.
func (p *WorkerPool) SubmitWait(task func()) {
	if task == nil {
		return
	}
	doneChan := make(chan struct{})
	p.taskQueue <- func() {
		task()
		close(doneChan)
	}
	<-doneChan
}

// dispatch sends the next queued task to an available worker.
func (p *WorkerPool) dispatch() {
	defer close(p.stoppedChan)
	timeout := time.NewTimer(p.timeout)
	var (
		workerCount    int
		task           func()
		ok, wait       bool
		workerTaskChan chan func()
	)
	startReady := make(chan chan func())
Loop:
	for {
		// As long as tasks are in the waiting queue, remove and execute these
		// tasks as workers become available, and place new incoming tasks on
		// the queue.  Once the queue is empty, then go back to submitting
		// incoming tasks directly to available workers.
		if p.waitingQueue.Len() != 0 {
			select {
			case task, ok = <-p.taskQueue:
				if !ok {
					break Loop
				}
				if task == nil {
					wait = true
					break Loop
				}
				p.waitingQueue.PushBack(task)
			case workerTaskChan = <-p.readyWorkers:
				// A worker is ready, so give task to worker.
				workerTaskChan <- p.waitingQueue.PopFront().(func())
			}
			continue
		}
		timeout.Reset(p.timeout)
		select {
		case task, ok = <-p.taskQueue:
			if !ok || task == nil {
				break Loop
			}
			// Got a task to do.
			select {
			case workerTaskChan = <-p.readyWorkers:
				// A worker is ready, so give task to worker.
				workerTaskChan <- task
			default:
				// No workers ready.
				// Create a new worker, if not at max.
				if workerCount < p.maxWorkers {
					workerCount++
					go func(t func()) {
						startWorker(startReady, p.readyWorkers)
						// Submit the task when the new worker.
						taskChan := <-startReady
						taskChan <- t
					}(task)
				} else {
					// Enqueue task to be executed by next available worker.
					p.waitingQueue.PushBack(task)
				}
			}
		case <-timeout.C:
			// Timed out waiting for work to arrive.  Kill a ready worker.
			if workerCount > 0 {
				select {
				case workerTaskChan = <-p.readyWorkers:
					// A worker is ready, so kill.
					close(workerTaskChan)
					workerCount--
				default:
					// No work, but no ready workers.  All workers are busy.
				}
			}
		}
	}

	// If instructed to wait for all queued tasks, then remove from queue and
	// give to workers until queue is empty.
	if wait {
		for p.waitingQueue.Len() != 0 {
			workerTaskChan = <-p.readyWorkers
			// A worker is ready, so give task to worker.
			workerTaskChan <- p.waitingQueue.PopFront().(func())
		}
	}

	// Stop all remaining workers as they become ready.
	for workerCount > 0 {
		workerTaskChan = <-p.readyWorkers
		close(workerTaskChan)
		workerCount--
	}
}

// startWorker starts a goroutine that executes tasks given by the dispatcher.
//
// When a new worker starts, it registers its availability on the startReady
// channel.  This ensures that the goroutine associated with starting the
// worker gets to use the worker to execute its task.  Otherwise, the main
// dispatcher loop could steal the new worker and not know to start up another
// worker for the waiting goroutine.  The task would then have to wait for
// another existing worker to become available, even though capacity is
// available to start additional workers.
//
// A worker registers that is it available to do work by putting its task
// channel on the readyWorkers channel.  The dispatcher reads a worker's task
// channel from the readyWorkers channel, and writes a task to the worker over
// the worker's task channel.  To stop a worker, the dispatcher closes a
// worker's task channel, instead of writing a task to it.
func startWorker(startReady, readyWorkers chan chan func()) {
	go func() {
		taskChan := make(chan func())
		var task func()
		var ok bool
		// Register availability on starReady channel.
		startReady <- taskChan
		for {
			// Read task from dispatcher.
			task, ok = <-taskChan
			if !ok {
				// Dispatcher has told worker to stop.
				break
			}

			// Execute the task.
			task()

			// Register availability on readyWorkers channel.
			readyWorkers <- taskChan
		}
	}()
}

// stop tells the dispatcher to exit, and whether or not to complete queued
// tasks.
func (p *WorkerPool) stop(wait bool) {
	p.stopMutex.Lock()
	defer p.stopMutex.Unlock()
	if p.stopped {
		return
	}
	p.stopped = true
	if wait {
		p.taskQueue <- nil
	}
	// Close task queue and wait for currently running tasks to finish.
	close(p.taskQueue)
	<-p.stoppedChan
}

// WaitingQueueSize will return the size of the waiting queue
func (p *WorkerPool) WaitingQueueSize() int {
	return p.waitingQueue.Len()
}

// TaksQueueSize will return the size of task queue
func (p *WorkerPool) TaksQueueSize() int {
	return len(p.taskQueue)
}

// QueueSize will return the size of the task queue first and waiting queue second
func (p *WorkerPool) QueueSize() (int, int) {
	return len(p.taskQueue), p.waitingQueue.Len()
}
<|endoftext|>


<commit_before>// The following enables go generate to generate the doc.go file.
//
//go:generate ./gendoc.sh

package main

import (
	"tools/presubmitter/impl"
)

func main() {
	impl.Root().Main()
}
<commit_msg>TBR: test<commit_after>// The following enables go generate to generate the doc.go file.
//
//go:generate ./gendoc.sh
//

package main

import (
	"tools/presubmitter/impl"
)

func main() {
	impl.Root().Main()
}
<|endoftext|>


<commit_before>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"encoding/binary"

	"github.com/google/gopacket"
)

// GRE is a Generic Routing Encapsulation header.
type GRE struct {
	BaseLayer
	ChecksumPresent, RoutingPresent, KeyPresent, SeqPresent, StrictSourceRoute bool
	RecursionControl, Flags, Version                                           uint8
	Protocol                                                                   EthernetType
	Checksum, Offset                                                           uint16
	Key, Seq                                                                   uint32
	*GRERouting
}

// GRERouting is GRE routing information, present if the RoutingPresent flag is
// set.
type GRERouting struct {
	AddressFamily        uint16
	SREOffset, SRELength uint8
	RoutingInformation   []byte
	Next                 *GRERouting
}

// LayerType returns gopacket.LayerTypeGRE.
func (g *GRE) LayerType() gopacket.LayerType { return LayerTypeGRE }

// DecodeFromBytes decodes the given bytes into this layer.
func (g *GRE) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	g.ChecksumPresent = data[0]&0x80 != 0
	g.RoutingPresent = data[0]&0x40 != 0
	g.KeyPresent = data[0]&0x20 != 0
	g.SeqPresent = data[0]&0x10 != 0
	g.StrictSourceRoute = data[0]&0x08 != 0
	g.RecursionControl = data[0] & 0x7
	g.Flags = data[1] >> 3
	g.Version = data[1] & 0x7
	g.Protocol = EthernetType(binary.BigEndian.Uint16(data[2:4]))
	offset := 4
	if g.ChecksumPresent || g.RoutingPresent {
		g.Checksum = binary.BigEndian.Uint16(data[offset : offset+2])
		g.Offset = binary.BigEndian.Uint16(data[offset+2 : offset+4])
		offset += 4
	}
	if g.KeyPresent {
		g.Key = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4
	}
	if g.SeqPresent {
		g.Seq = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4
	}
	if g.RoutingPresent {
		tail := &g.GRERouting
		for {
			sre := &GRERouting{
				AddressFamily: binary.BigEndian.Uint16(data[offset : offset+2]),
				SREOffset:     data[offset+2],
				SRELength:     data[offset+3],
			}
			sre.RoutingInformation = data[offset+4 : offset+4+int(sre.SRELength)]
			offset += 4 + int(sre.SRELength)
			if sre.AddressFamily == 0 && sre.SRELength == 0 {
				break
			}
			(*tail) = sre
			tail = &sre.Next
		}
	}
	g.BaseLayer = BaseLayer{data[:offset], data[offset:]}
	return nil
}

// SerializeTo writes the serialized form of this layer into the SerializationBuffer,
// implementing gopacket.SerializableLayer. See the docs for gopacket.SerializableLayer for more info.
func (g *GRE) SerializeTo(b gopacket.SerializeBuffer, opts gopacket.SerializeOptions) error {
	size := 4
	if g.ChecksumPresent || g.RoutingPresent {
		size += 4
	}
	if g.KeyPresent {
		size += 4
	}
	if g.SeqPresent {
		size += 4
	}
	if g.RoutingPresent {
		r := g.GRERouting
		for r != nil {
			size += 4 + int(r.SRELength)
			r = r.Next
		}
		size += 4
	}
	buf, err := b.PrependBytes(size)
	if err != nil {
		return err
	}
	// Reset any potentially dirty memory in the first 2 bytes, as these use OR to set flags.
	buf[0] = 0
	buf[1] = 0
	if g.ChecksumPresent {
		buf[0] |= 0x80
	}
	if g.RoutingPresent {
		buf[0] |= 0x40
	}
	if g.KeyPresent {
		buf[0] |= 0x20
	}
	if g.SeqPresent {
		buf[0] |= 0x10
	}
	if g.StrictSourceRoute {
		buf[0] |= 0x08
	}
	buf[0] |= g.RecursionControl
	buf[1] |= g.Flags << 3
	buf[1] |= g.Version
	binary.BigEndian.PutUint16(buf[2:4], uint16(g.Protocol))
	offset := 4
	if g.ChecksumPresent || g.RoutingPresent {
		// Don't write the checksum value yet, as we may need to compute it,
		// which requires the entire header be complete.
		// Instead we zeroize the memory in case it is dirty.
		buf[offset] = 0
		buf[offset+1] = 0
		binary.BigEndian.PutUint16(buf[offset+2:offset+4], g.Offset)
		offset += 4
	}
	if g.KeyPresent {
		binary.BigEndian.PutUint32(buf[offset:offset+4], g.Key)
		offset += 4
	}
	if g.SeqPresent {
		binary.BigEndian.PutUint32(buf[offset:offset+4], g.Seq)
		offset += 4
	}
	if g.RoutingPresent {
		sre := g.GRERouting
		for sre != nil {
			binary.BigEndian.PutUint16(buf[offset:offset+2], sre.AddressFamily)
			buf[offset+2] = sre.SREOffset
			buf[offset+3] = sre.SRELength
			copy(buf[offset+4:offset+4+int(sre.SRELength)], sre.RoutingInformation)
			offset += 4 + int(sre.SRELength)
			sre = sre.Next
		}
		// Terminate routing field with a "NULL" SRE.
		binary.BigEndian.PutUint32(buf[offset:offset+4], 0)
	}
	if g.ChecksumPresent {
		if opts.ComputeChecksums {
			g.Checksum = tcpipChecksum(b.Bytes(), 0)
		}

		binary.BigEndian.PutUint16(buf[4:6], g.Checksum)
	}
	return nil
}

// CanDecode returns the set of layer types that this DecodingLayer can decode.
func (g *GRE) CanDecode() gopacket.LayerClass {
	return LayerTypeGRE
}

// NextLayerType returns the layer type contained by this DecodingLayer.
func (g *GRE) NextLayerType() gopacket.LayerType {
	return g.Protocol.LayerType()
}

func decodeGRE(data []byte, p gopacket.PacketBuilder) error {
	g := &GRE{}
	return decodingLayerDecoder(g, data, p)
}
<commit_msg>Add support for GRE Ack as used in PPTP<commit_after>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"encoding/binary"

	"github.com/google/gopacket"
)

// GRE is a Generic Routing Encapsulation header.
type GRE struct {
	BaseLayer
	ChecksumPresent, RoutingPresent, KeyPresent, SeqPresent, StrictSourceRoute, AckPresent bool
	RecursionControl, Flags, Version                                                       uint8
	Protocol                                                                               EthernetType
	Checksum, Offset                                                                       uint16
	Key, Seq, Ack                                                                          uint32
	*GRERouting
}

// GRERouting is GRE routing information, present if the RoutingPresent flag is
// set.
type GRERouting struct {
	AddressFamily        uint16
	SREOffset, SRELength uint8
	RoutingInformation   []byte
	Next                 *GRERouting
}

// LayerType returns gopacket.LayerTypeGRE.
func (g *GRE) LayerType() gopacket.LayerType { return LayerTypeGRE }

// DecodeFromBytes decodes the given bytes into this layer.
func (g *GRE) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	g.ChecksumPresent = data[0]&0x80 != 0
	g.RoutingPresent = data[0]&0x40 != 0
	g.KeyPresent = data[0]&0x20 != 0
	g.SeqPresent = data[0]&0x10 != 0
	g.StrictSourceRoute = data[0]&0x08 != 0
	g.AckPresent = data[1]&0x80 != 0
	g.RecursionControl = data[0] & 0x7
	g.Flags = data[1] >> 3
	g.Version = data[1] & 0x7
	g.Protocol = EthernetType(binary.BigEndian.Uint16(data[2:4]))
	offset := 4
	if g.ChecksumPresent || g.RoutingPresent {
		g.Checksum = binary.BigEndian.Uint16(data[offset : offset+2])
		g.Offset = binary.BigEndian.Uint16(data[offset+2 : offset+4])
		offset += 4
	}
	if g.KeyPresent {
		g.Key = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4
	}
	if g.SeqPresent {
		g.Seq = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4
	}
	if g.RoutingPresent {
		tail := &g.GRERouting
		for {
			sre := &GRERouting{
				AddressFamily: binary.BigEndian.Uint16(data[offset : offset+2]),
				SREOffset:     data[offset+2],
				SRELength:     data[offset+3],
			}
			sre.RoutingInformation = data[offset+4 : offset+4+int(sre.SRELength)]
			offset += 4 + int(sre.SRELength)
			if sre.AddressFamily == 0 && sre.SRELength == 0 {
				break
			}
			(*tail) = sre
			tail = &sre.Next
		}
	}
	if g.AckPresent {
		g.Ack = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4
	}
	g.BaseLayer = BaseLayer{data[:offset], data[offset:]}
	return nil
}

// SerializeTo writes the serialized form of this layer into the SerializationBuffer,
// implementing gopacket.SerializableLayer. See the docs for gopacket.SerializableLayer for more info.
func (g *GRE) SerializeTo(b gopacket.SerializeBuffer, opts gopacket.SerializeOptions) error {
	size := 4
	if g.ChecksumPresent || g.RoutingPresent {
		size += 4
	}
	if g.KeyPresent {
		size += 4
	}
	if g.SeqPresent {
		size += 4
	}
	if g.RoutingPresent {
		r := g.GRERouting
		for r != nil {
			size += 4 + int(r.SRELength)
			r = r.Next
		}
		size += 4
	}
	if g.AckPresent {
		size += 4
	}
	buf, err := b.PrependBytes(size)
	if err != nil {
		return err
	}
	// Reset any potentially dirty memory in the first 2 bytes, as these use OR to set flags.
	buf[0] = 0
	buf[1] = 0
	if g.ChecksumPresent {
		buf[0] |= 0x80
	}
	if g.RoutingPresent {
		buf[0] |= 0x40
	}
	if g.KeyPresent {
		buf[0] |= 0x20
	}
	if g.SeqPresent {
		buf[0] |= 0x10
	}
	if g.StrictSourceRoute {
		buf[0] |= 0x08
	}
	if g.AckPresent {
		buf[1] |= 0x80
	}
	buf[0] |= g.RecursionControl
	buf[1] |= g.Flags << 3
	buf[1] |= g.Version
	binary.BigEndian.PutUint16(buf[2:4], uint16(g.Protocol))
	offset := 4
	if g.ChecksumPresent || g.RoutingPresent {
		// Don't write the checksum value yet, as we may need to compute it,
		// which requires the entire header be complete.
		// Instead we zeroize the memory in case it is dirty.
		buf[offset] = 0
		buf[offset+1] = 0
		binary.BigEndian.PutUint16(buf[offset+2:offset+4], g.Offset)
		offset += 4
	}
	if g.KeyPresent {
		binary.BigEndian.PutUint32(buf[offset:offset+4], g.Key)
		offset += 4
	}
	if g.SeqPresent {
		binary.BigEndian.PutUint32(buf[offset:offset+4], g.Seq)
		offset += 4
	}
	if g.RoutingPresent {
		sre := g.GRERouting
		for sre != nil {
			binary.BigEndian.PutUint16(buf[offset:offset+2], sre.AddressFamily)
			buf[offset+2] = sre.SREOffset
			buf[offset+3] = sre.SRELength
			copy(buf[offset+4:offset+4+int(sre.SRELength)], sre.RoutingInformation)
			offset += 4 + int(sre.SRELength)
			sre = sre.Next
		}
		// Terminate routing field with a "NULL" SRE.
		binary.BigEndian.PutUint32(buf[offset:offset+4], 0)
	}
	if g.AckPresent {
		binary.BigEndian.PutUint32(buf[offset:offset+4], g.Ack)
		offset += 4
	}
	if g.ChecksumPresent {
		if opts.ComputeChecksums {
			g.Checksum = tcpipChecksum(b.Bytes(), 0)
		}

		binary.BigEndian.PutUint16(buf[4:6], g.Checksum)
	}
	return nil
}

// CanDecode returns the set of layer types that this DecodingLayer can decode.
func (g *GRE) CanDecode() gopacket.LayerClass {
	return LayerTypeGRE
}

// NextLayerType returns the layer type contained by this DecodingLayer.
func (g *GRE) NextLayerType() gopacket.LayerType {
	return g.Protocol.LayerType()
}

func decodeGRE(data []byte, p gopacket.PacketBuilder) error {
	g := &GRE{}
	return decodingLayerDecoder(g, data, p)
}
<|endoftext|>


<commit_before>/*
Copyright 2011 Google Inc.

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

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

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

package search

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"time"

	"camli/blobref"
)

type Result struct {
	BlobRef     *blobref.BlobRef
	Signer      *blobref.BlobRef // may be nil
	LastModTime int64            // seconds since epoch
}

// Results exists mostly for debugging, to provide a String method on
// a slice of Result.
type Results []*Result

func (s Results) String() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "[%d search results: ", len(s))
	for _, r := range s {
		fmt.Fprintf(&buf, "{BlobRef: %s, Signer: %s, LastModTime: %d}",
			r.BlobRef, r.Signer, r.LastModTime)
	}
	buf.WriteString("]")
	return buf.String()
}

// TODO: move this to schema or something?
type Claim struct {
	BlobRef, Signer, Permanode *blobref.BlobRef

	Date *time.Time
	Type string // "set-attribute", "add-attribute", etc

	// If an attribute modification
	Attr, Value string
}

func (c *Claim) String() string {
	return fmt.Sprintf(
		"search.Claim{BlobRef: %s, Signer: %s, Permanode: %s, Date: %s, Type: %s, Attr: %s, Value: %s}",
		c.BlobRef, c.Signer, c.Permanode, c.Date, c.Type, c.Attr, c.Value)
}

type ClaimList []*Claim

func (cl ClaimList) Len() int {
	return len(cl)
}

func (cl ClaimList) Less(i, j int) bool {
	// TODO: memoize Seconds in unexported Claim field
	return cl[i].Date.Seconds() < cl[j].Date.Seconds()
}

func (cl ClaimList) Swap(i, j int) {
	cl[i], cl[j] = cl[j], cl[i]
}

func (cl ClaimList) String() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "[%d claims: ", len(cl))
	for _, r := range cl {
		buf.WriteString(r.String())
	}
	buf.WriteString("]")
	return buf.String()
}

type FileInfo struct {
	Size     int64  `json:"size"`
	FileName string `json:"fileName"`
	MimeType string `json:"mimeType"`
}

func (fi *FileInfo) IsImage() bool {
	return strings.HasPrefix(fi.MimeType, "image/")
}

type Path struct {
	Claim, Base, Target *blobref.BlobRef
	ClaimDate           string
	Suffix              string
}

func (p *Path) String() string {
	return fmt.Sprintf("Path{Claim: %v, %v; Base: %v + Suffix %q => Target %v}",
		p.Claim, p.ClaimDate, p.Base, p.Suffix, p.Target)
}

type PermanodeByAttrRequest struct {
	Attribute  string // currently supported: "tag", "title"
	Query      string
	Signer     *blobref.BlobRef
	FuzzyMatch bool // by default, an exact match is required
	MaxResults int  // optional max results
}

type Index interface {
	// dest must be closed, even when returning an error.
	// limit is <= 0 for default.  smallest possible default is 0
	GetRecentPermanodes(dest chan *Result,
		owner *blobref.BlobRef,
		limit int) os.Error

	// SearchPermanodes finds permanodes matching the provided
	// request and sends unique permanode blobrefs to dest.
	// In particular, if request.FuzzyMatch is true, a fulltext
	// search is performed (if supported by the attribute(s))
	// instead of an exact match search.
	// Additionally, if request.Attribute is blank, all attributes
	// are searched (as fulltext), otherwise the search is
	// restricted  to the named attribute.
	//
	// dest is always closed, regardless of the error return value.
	SearchPermanodesWithAttr(dest chan<- *blobref.BlobRef,
		request *PermanodeByAttrRequest) os.Error

	GetOwnerClaims(permaNode, owner *blobref.BlobRef) (ClaimList, os.Error)

	// os.ENOENT should be returned if the blob isn't known
	GetBlobMimeType(blob *blobref.BlobRef) (mime string, size int64, err os.Error)

	// ExistingFileSchemas returns 0 or more blobrefs of file
	// schema blobs that represent the bytes of a file given in
	// bytesRef.  The file schema blobs returned are not
	// guaranteed to reference chunks that still exist on the
	// blobservers, though.  It's purely a hint for clients to
	// avoid uploads if possible.  Before re-using any returned
	// blobref they should be checked.
	ExistingFileSchemas(bytesRef *blobref.BlobRef) ([]*blobref.BlobRef, os.Error)

	GetFileInfo(fileRef *blobref.BlobRef) (*FileInfo, os.Error)

	// Given an owner key, a camliType 'claim', 'attribute' name,
	// and specific 'value', find the most recent permanode that has
	// a corresponding 'set-attribute' claim attached.
	// Returns os.ENOENT if none is found.
	// Only attributes white-listed by IsIndexedAttribute are valid.
	PermanodeOfSignerAttrValue(signer *blobref.BlobRef, attr, val string) (*blobref.BlobRef, os.Error)

	// PathsOfSignerTarget queries the index about "camliPath:"
	// URL-dispatch attributes.
	//
	// It returns a list of all the path claims that have been signed
	// by the provided signer and point at the given target.
	//
	// This is used when editing a permanode, to figure work up
	// the name resolution tree backwards ultimately to a
	// camliRoot permanode (which should know its base URL), and
	// then the complete URL(s) of a target can be found.
	PathsOfSignerTarget(signer, target *blobref.BlobRef) ([]*Path, os.Error)

	// All Path claims for (signer, base, suffix)
	PathsLookup(signer, base *blobref.BlobRef, suffix string) ([]*Path, os.Error)

	// Most recent Path claim for (signer, base, suffix) as of
	// provided time 'at', or most recent if 'at' is nil.
	PathLookup(signer, base *blobref.BlobRef, suffix string, at *time.Time) (*Path, os.Error)
}

// TODO(bradfitz): rename this? This is really about signer-attr-value
// (PermanodeOfSignerAttrValue), and not about indexed attributes in general.
func IsIndexedAttribute(attr string) bool {
	switch attr {
	case "camliRoot", "tag", "title":
		return true
	}
	return false
}

func IsFulltextAttribute(attr string) bool {
	switch attr {
	case "tag", "title":
		return true
	}
	return false
}
<commit_msg>index: bigger comment on ExistingFileSchemas<commit_after>/*
Copyright 2011 Google Inc.

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

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

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

package search

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"time"

	"camli/blobref"
)

type Result struct {
	BlobRef     *blobref.BlobRef
	Signer      *blobref.BlobRef // may be nil
	LastModTime int64            // seconds since epoch
}

// Results exists mostly for debugging, to provide a String method on
// a slice of Result.
type Results []*Result

func (s Results) String() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "[%d search results: ", len(s))
	for _, r := range s {
		fmt.Fprintf(&buf, "{BlobRef: %s, Signer: %s, LastModTime: %d}",
			r.BlobRef, r.Signer, r.LastModTime)
	}
	buf.WriteString("]")
	return buf.String()
}

// TODO: move this to schema or something?
type Claim struct {
	BlobRef, Signer, Permanode *blobref.BlobRef

	Date *time.Time
	Type string // "set-attribute", "add-attribute", etc

	// If an attribute modification
	Attr, Value string
}

func (c *Claim) String() string {
	return fmt.Sprintf(
		"search.Claim{BlobRef: %s, Signer: %s, Permanode: %s, Date: %s, Type: %s, Attr: %s, Value: %s}",
		c.BlobRef, c.Signer, c.Permanode, c.Date, c.Type, c.Attr, c.Value)
}

type ClaimList []*Claim

func (cl ClaimList) Len() int {
	return len(cl)
}

func (cl ClaimList) Less(i, j int) bool {
	// TODO: memoize Seconds in unexported Claim field
	return cl[i].Date.Seconds() < cl[j].Date.Seconds()
}

func (cl ClaimList) Swap(i, j int) {
	cl[i], cl[j] = cl[j], cl[i]
}

func (cl ClaimList) String() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "[%d claims: ", len(cl))
	for _, r := range cl {
		buf.WriteString(r.String())
	}
	buf.WriteString("]")
	return buf.String()
}

type FileInfo struct {
	Size     int64  `json:"size"`
	FileName string `json:"fileName"`
	MimeType string `json:"mimeType"`
}

func (fi *FileInfo) IsImage() bool {
	return strings.HasPrefix(fi.MimeType, "image/")
}

type Path struct {
	Claim, Base, Target *blobref.BlobRef
	ClaimDate           string
	Suffix              string
}

func (p *Path) String() string {
	return fmt.Sprintf("Path{Claim: %v, %v; Base: %v + Suffix %q => Target %v}",
		p.Claim, p.ClaimDate, p.Base, p.Suffix, p.Target)
}

type PermanodeByAttrRequest struct {
	Attribute  string // currently supported: "tag", "title"
	Query      string
	Signer     *blobref.BlobRef
	FuzzyMatch bool // by default, an exact match is required
	MaxResults int  // optional max results
}

type Index interface {
	// dest must be closed, even when returning an error.
	// limit is <= 0 for default.  smallest possible default is 0
	GetRecentPermanodes(dest chan *Result,
		owner *blobref.BlobRef,
		limit int) os.Error

	// SearchPermanodes finds permanodes matching the provided
	// request and sends unique permanode blobrefs to dest.
	// In particular, if request.FuzzyMatch is true, a fulltext
	// search is performed (if supported by the attribute(s))
	// instead of an exact match search.
	// Additionally, if request.Attribute is blank, all attributes
	// are searched (as fulltext), otherwise the search is
	// restricted  to the named attribute.
	//
	// dest is always closed, regardless of the error return value.
	SearchPermanodesWithAttr(dest chan<- *blobref.BlobRef,
		request *PermanodeByAttrRequest) os.Error

	GetOwnerClaims(permaNode, owner *blobref.BlobRef) (ClaimList, os.Error)

	// os.ENOENT should be returned if the blob isn't known
	GetBlobMimeType(blob *blobref.BlobRef) (mime string, size int64, err os.Error)

	// ExistingFileSchemas returns 0 or more blobrefs of "bytes"
	// (TODO(bradfitz): or file?) schema blobs that represent the
	// bytes of a file given in bytesRef.  The file schema blobs
	// returned are not guaranteed to reference chunks that still
	// exist on the blobservers, though.  It's purely a hint for
	// clients to avoid uploads if possible.  Before re-using any
	// returned blobref they should be checked.
	//
	// Use case: a user drag & drops a large file onto their
	// browser to upload.  (imagine that "large" means anything
	// larger than a blobserver's max blob size) JavaScript can
	// first SHA-1 the large file locally, then send the
	// wholeFileRef to this call and see if they'd previously
	// uploaded the same file in the past.  If so, the upload
	// can be avoided if at least one of the returned schemaRefs
	// can be validated (with a validating HEAD request) to still
	// all exist on the blob server.
	ExistingFileSchemas(wholeFileRef *blobref.BlobRef) (schemaRefs []*blobref.BlobRef, os.Error)

	GetFileInfo(fileRef *blobref.BlobRef) (*FileInfo, os.Error)

	// Given an owner key, a camliType 'claim', 'attribute' name,
	// and specific 'value', find the most recent permanode that has
	// a corresponding 'set-attribute' claim attached.
	// Returns os.ENOENT if none is found.
	// Only attributes white-listed by IsIndexedAttribute are valid.
	PermanodeOfSignerAttrValue(signer *blobref.BlobRef, attr, val string) (*blobref.BlobRef, os.Error)

	// PathsOfSignerTarget queries the index about "camliPath:"
	// URL-dispatch attributes.
	//
	// It returns a list of all the path claims that have been signed
	// by the provided signer and point at the given target.
	//
	// This is used when editing a permanode, to figure work up
	// the name resolution tree backwards ultimately to a
	// camliRoot permanode (which should know its base URL), and
	// then the complete URL(s) of a target can be found.
	PathsOfSignerTarget(signer, target *blobref.BlobRef) ([]*Path, os.Error)

	// All Path claims for (signer, base, suffix)
	PathsLookup(signer, base *blobref.BlobRef, suffix string) ([]*Path, os.Error)

	// Most recent Path claim for (signer, base, suffix) as of
	// provided time 'at', or most recent if 'at' is nil.
	PathLookup(signer, base *blobref.BlobRef, suffix string, at *time.Time) (*Path, os.Error)
}

// TODO(bradfitz): rename this? This is really about signer-attr-value
// (PermanodeOfSignerAttrValue), and not about indexed attributes in general.
func IsIndexedAttribute(attr string) bool {
	switch attr {
	case "camliRoot", "tag", "title":
		return true
	}
	return false
}

func IsFulltextAttribute(attr string) bool {
	switch attr {
	case "tag", "title":
		return true
	}
	return false
}
<|endoftext|>


<commit_before>/*
Package text holds models which
make text classification easy. They
are regular models, but take strings
as arguments so you can feed in
documents rather than large,
hand-constructed word vectors. Although
models might represent the words as
these vectors, the munging of a
document is hidden from the user.

The simplest model, although suprisingly
effective, is Naive Bayes. If you
want to read more about the specific
model, check out the docs for the
NaiveBayes struct/model.

The following example is an online Naive
Bayes model used for sentiment analysis.

Example Online Naive Bayes Text Classifier (multiclass):

	// create the channel of data and errors
	stream := make(chan base.TextDatapoint, 100)
	errors := make(chan error)

	// make a new NaiveBayes model with
	// 2 classes expected (classes in
	// datapoints will now expect {0,1}.
	// in general, given n as the classes
	// variable, the model will expect
	// datapoint classes in {0,...,n-1})
	//
	// Note that the model is filtering
	// the text to omit anything except
	// words and numbers (and spaces
	// obviously)
	model := NewNaiveBayes(stream, 2, base.OnlyWordsAndNumbers)

	go model.OnlineLearn(errors)

	stream <- base.TextDatapoint{
		X: "I love the city",
		Y: 1,
	}

	stream <- base.TextDatapoint{
		X: "I hate Los Angeles",
		Y: 0,
	}

	stream <- base.TextDatapoint{
		X: "My mother is not a nice lady",
		Y: 0,
	}

	close(stream)

	for {
		err, more := <- errors
		if err != nil {
			fmt.Fprintf(b.Output, "Error passed: %v", err)
		} else {
			// training is done!
			break
		}
	}

	// now you can predict like normal
	class := model.Predict("My mother is in Los Angeles") // 0
*/
package text

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"strings"
	"sync"

	"golang.org/x/text/transform"

	"github.com/piazzamp/goml/base"
)

/*
NaiveBayes is a general classification
model that calculates the probability
that a datapoint is part of a class
by using Bayes Rule:
	P(y|x) = P(x|y)*P(y)/P(x)
The unique part of this model is that
it assumes words are unrelated to
eachother. For example, the probability
of seeing the word 'penis' in spam emails
if you've already seen 'viagra' might be
different than if you hadn't seen it. The
model ignores this fact because the
computation of full Bayesian model would
take much longer, and would grow significantly
with each word you see.

https://en.wikipedia.org/wiki/Naive_Bayes_classifier
http://cs229.stanford.edu/notes/cs229-notes2.pdf

Based on Bayes Rule, we can easily calculate
the numerator (x | y is just the number of
times x is seen and the class=y, and P(y) is
just the number of times y=class / the number
of positive training examples/words.) The
denominator is also easy to calculate, but
if you recognize that it's just a constant
because it's just the probability of seeing
a certain document given the dataset we can
make the following transformation to be
able to classify without as much classification:
	Class(x) = argmax_c{P(y = c) * ∏P(x|y = c)}
And we can use logarithmic transformations to
make this calculation more computer-practical
(multiplying a bunch of probabilities on [0,1]
will always result in a very small number
which could easily underflow the float value):
	Class(x) = argmax_c{log(P(y = c)) + ΣP(x|y = c)}
Much better. That's our model!
*/
type NaiveBayes struct {
	// Words holds a map of words
	// to their corresponding Word
	// structure
	Words histogram `json:"words"`

	// Count holds the number of times
	// class i was seen as Count[i]
	Count []uint64 `json:"count"`

	// Probabilities holds the probability
	// that class Y is class i as
	// Probabilities[i] for
	Probabilities []float64 `json:"probabilities"`

	// DocumentCount holds the number of
	// documents that have been seen
	DocumentCount uint64 `json:"document_count"`

	// DictCount holds the size of the
	// NaiveBayes model's vocabulary
	DictCount uint64 `json:"vocabulary_size"`

	// sanitize is used by a model
	// to sanitize input of text
	sanitize transform.Transformer

	// stream holds the datastream
	stream <-chan base.TextDatapoint

	// Output is the io.Writer used for logging
	// and printing. Defaults to os.Stdout.
	Output io.Writer
}

// histogram allows conncurrency-friendly map access via its
// exported Get and Set methods
type histogram struct {
	sync.RWMutex
	words map[string]Word
}

// Get looks up a word from h's Word map, it should be used in
// place of a direct map lookup
// the only caveat here is that it will always return the 'success' boolean
func (h *histogram) Get(w string) (Word, bool) {
	h.RLock()
	result, ok := h.words[w]
	h.RUnlock()
	return result, ok
}

// Set sets word k's value to v in h's Word map
func (h *histogram) Set(k string, v Word) {
	h.Lock()
	h.words[k] = v
	h.Unlock()
}

// Word holds the structural
// information needed to calculate
// the probability of
type Word struct {
	// Count holds the number of times,
	// (i in Count[i] is the given class)
	Count []uint64

	// Seen holds the number of times
	// the world has been seen. This
	// is than same as
	//    foldl (+) 0 Count
	// in Haskell syntax, but is included
	// you wouldn't have to calculate
	// this every time you wanted to
	// recalc the probabilities (foldl
	// is the same as reduce, basically.)
	Seen uint64

	// DocsSeen is the same as Seen but
	// a word is only counted once even
	// if it's in a document multiple times
	DocsSeen uint64
}

// NewNaiveBayes returns a NaiveBayes model the
// given number of classes instantiated, ready
// to learn off the given data stream. The sanitization
// function is set to the given function. It must
// comply with the transform.RemoveFunc interface
func NewNaiveBayes(stream <-chan base.TextDatapoint, classes uint8, sanitize func(rune) bool) *NaiveBayes {
	return &NaiveBayes{
		Words:         histogram{sync.RWMutex{}, make(map[string]Word)},
		Count:         make([]uint64, classes),
		Probabilities: make([]float64, classes),

		sanitize: transform.RemoveFunc(sanitize),
		stream:   stream,

		Output: os.Stdout,
	}
}

// Predict takes in a document, predicts the
// class of the document based on the training
// data passed so far, and returns the class
// estimated for the document.
func (b *NaiveBayes) Predict(sentence string) uint8 {
	sums := make([]float64, len(b.Count))

	sentence, _, _ = transform.String(b.sanitize, sentence)
	words := strings.Split(strings.ToLower(sentence), " ")
	for _, word := range words {
		w, ok := b.Words.Get(word)
		if !ok {
			continue
		}

		for i := range sums {
			sums[i] += math.Log(float64(w.Count[i]+1) / float64(w.Seen+b.DictCount))
		}
	}

	for i := range sums {
		sums[i] += math.Log(b.Probabilities[i])
	}

	// find best class
	var maxI int
	for i := range sums {
		if sums[i] > sums[maxI] {
			maxI = i
		}
	}

	return uint8(maxI)
}

// Probability takes in a small document, returns the
// estimated class of the document based on the model
// as well as the probability that the model is part
// of that class
//
// NOTE: you should only use this for small documents
// because, as discussed in the docs for the model, the
// probability will often times underflow because you
// are multiplying together a bunch of probabilities
// which range on [0,1]. As such, the returned float
// could be NaN, and the predicted class could be
// 0 always.
//
// Basically, use Predict to be robust for larger
// documents. Use Probability only on relatively small
// (MAX of maybe a dozen words - basically just
// sentences and words) documents.
func (b *NaiveBayes) Probability(sentence string) (uint8, float64) {
	sums := make([]float64, len(b.Count))
	for i := range sums {
		sums[i] = 1
	}

	sentence, _, _ = transform.String(b.sanitize, sentence)
	words := strings.Split(strings.ToLower(sentence), " ")
	for _, word := range words {
		w, ok := b.Words.Get(word)
		if !ok {
			continue
		}

		for i := range sums {
			sums[i] *= float64(w.Count[i]+1) / float64(w.Seen+b.DictCount)
		}
	}

	for i := range sums {
		sums[i] *= b.Probabilities[i]
	}

	var denom float64
	var maxI int
	for i := range sums {
		if sums[i] > sums[maxI] {
			maxI = i
		}

		denom += sums[i]
	}

	return uint8(maxI), sums[maxI] / denom
}

// OnlineLearn lets the NaiveBayes model learn
// from the datastream, waiting for new data to
// come into the stream from a separate goroutine
func (b *NaiveBayes) OnlineLearn(errors chan<- error) {
	if errors == nil {
		errors = make(chan error)
	}
	if b.stream == nil {
		errors <- fmt.Errorf("ERROR: attempting to learn with nil data stream!\n")
		close(errors)
		return
	}

	fmt.Fprintf(b.Output, "Training:\n\tModel: Multinomial Naïve Bayes\n\tClasses: %v\n", len(b.Count))

	var point base.TextDatapoint
	var more bool

	for {
		point, more = <-b.stream

		if more {
			// sanitize and break up document
			sanitized, _, _ := transform.String(b.sanitize, point.X)
			sanitized = strings.ToLower(sanitized)

			words := strings.Split(sanitized, " ")

			C := int(point.Y)

			if C > len(b.Count)-1 {
				errors <- fmt.Errorf("ERROR: given document class is greater than the number of classes in the model!\n")
				continue
			}

			// update global class probabilities
			b.Count[C]++
			b.DocumentCount++
			for i := range b.Probabilities {
				b.Probabilities[i] = float64(b.Count[i]) / float64(b.DocumentCount)
			}

			// store words seen in document (to add to DocsSeen)
			seenCount := make(map[string]int)

			// update probabilities for words
			for _, word := range words {
				if len(word) < 3 {
					continue
				}

				w, ok := b.Words.Get(word)

				if !ok {
					w = Word{
						Count: make([]uint64, len(b.Count)),
						Seen:  uint64(0),
					}

					b.DictCount++
				}

				w.Count[C]++
				w.Seen++

				b.Words.Set(word, w)

				seenCount[word] = 1
			}

			// add to DocsSeen
			for term := range seenCount {
				tmp, _ := b.Words.Get(term)
				tmp.DocsSeen++
				b.Words.Set(term, tmp)
			}
		} else {
			fmt.Fprintf(b.Output, "Training Completed.\n%v\n\n", b)
			close(errors)
			return
		}
	}
}

// UpdateStream updates the NaiveBayes model's
// text datastream
func (b *NaiveBayes) UpdateStream(stream chan base.TextDatapoint) {
	b.stream = stream
}

// UpdateSanitize updates the NaiveBayes model's
// text sanitization transformation function
func (b *NaiveBayes) UpdateSanitize(sanitize func(rune) bool) {
	b.sanitize = transform.RemoveFunc(sanitize)
}

// String implements the fmt interface for clean printing. Here
// we're using it to print the model as the equation h(θ)=...
// where h is the perceptron hypothesis model.
func (b *NaiveBayes) String() string {
	return fmt.Sprintf("h(θ) = argmax_c{log(P(y = c)) + Σlog(P(x|y = c))}\n\tClasses: %v\n\tDocuments evaluated in model: %v\n\tWords evaluated in model: %v\n", len(b.Count), int(b.DocumentCount), int(b.DictCount))
}

// PersistToFile takes in an absolute filepath and saves the
// parameter vector θ to the file, which can be restored later.
// The function will take paths from the current directory, but
// functions
//
// The data is stored as JSON because it's one of the most
// efficient storage method (you only need one comma extra
// per feature + two brackets, total!) And it's extendable.
func (b *NaiveBayes) PersistToFile(path string) error {
	if path == "" {
		return fmt.Errorf("ERROR: you just tried to persist your model to a file with no path!! That's a no-no. Try it with a valid filepath")
	}

	bytes, err := json.Marshal(b)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(path, bytes, os.ModePerm)
	if err != nil {
		return err
	}

	return nil
}

// Restore takes the bytes of a NaiveBayes model and
// restores a model to it. This would be useful if
// training a model and saving it into a project
// using go-bindata (look it up) so you don't have
// to persist a large file and deal with paths on
// a production system. This option is included
// in text models vs. others because the text models
// usually have much larger storage requirements
func (b *NaiveBayes) Restore(bytes []byte) error {
	err := json.Unmarshal(bytes, &b)
	if err != nil {
		return err
	}

	return nil
}

// RestoreFromFile takes in a path to a parameter vector theta
// and assigns the model it's operating on's parameter vector
// to that.
//
// The path must ba an absolute path or a path from the current
// directory
//
// This would be useful in persisting data between running
// a model on data.
func (b *NaiveBayes) RestoreFromFile(path string) error {
	if path == "" {
		return fmt.Errorf("ERROR: you just tried to restore your model from a file with no path! That's a no-no. Try it with a valid filepath")
	}

	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	err = json.Unmarshal(bytes, &b)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>fix import path<commit_after>/*
Package text holds models which
make text classification easy. They
are regular models, but take strings
as arguments so you can feed in
documents rather than large,
hand-constructed word vectors. Although
models might represent the words as
these vectors, the munging of a
document is hidden from the user.

The simplest model, although suprisingly
effective, is Naive Bayes. If you
want to read more about the specific
model, check out the docs for the
NaiveBayes struct/model.

The following example is an online Naive
Bayes model used for sentiment analysis.

Example Online Naive Bayes Text Classifier (multiclass):

	// create the channel of data and errors
	stream := make(chan base.TextDatapoint, 100)
	errors := make(chan error)

	// make a new NaiveBayes model with
	// 2 classes expected (classes in
	// datapoints will now expect {0,1}.
	// in general, given n as the classes
	// variable, the model will expect
	// datapoint classes in {0,...,n-1})
	//
	// Note that the model is filtering
	// the text to omit anything except
	// words and numbers (and spaces
	// obviously)
	model := NewNaiveBayes(stream, 2, base.OnlyWordsAndNumbers)

	go model.OnlineLearn(errors)

	stream <- base.TextDatapoint{
		X: "I love the city",
		Y: 1,
	}

	stream <- base.TextDatapoint{
		X: "I hate Los Angeles",
		Y: 0,
	}

	stream <- base.TextDatapoint{
		X: "My mother is not a nice lady",
		Y: 0,
	}

	close(stream)

	for {
		err, more := <- errors
		if err != nil {
			fmt.Fprintf(b.Output, "Error passed: %v", err)
		} else {
			// training is done!
			break
		}
	}

	// now you can predict like normal
	class := model.Predict("My mother is in Los Angeles") // 0
*/
package text

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"strings"
	"sync"

	"golang.org/x/text/transform"

	"github.com/cdipaolo/goml/base"
)

/*
NaiveBayes is a general classification
model that calculates the probability
that a datapoint is part of a class
by using Bayes Rule:
	P(y|x) = P(x|y)*P(y)/P(x)
The unique part of this model is that
it assumes words are unrelated to
eachother. For example, the probability
of seeing the word 'penis' in spam emails
if you've already seen 'viagra' might be
different than if you hadn't seen it. The
model ignores this fact because the
computation of full Bayesian model would
take much longer, and would grow significantly
with each word you see.

https://en.wikipedia.org/wiki/Naive_Bayes_classifier
http://cs229.stanford.edu/notes/cs229-notes2.pdf

Based on Bayes Rule, we can easily calculate
the numerator (x | y is just the number of
times x is seen and the class=y, and P(y) is
just the number of times y=class / the number
of positive training examples/words.) The
denominator is also easy to calculate, but
if you recognize that it's just a constant
because it's just the probability of seeing
a certain document given the dataset we can
make the following transformation to be
able to classify without as much classification:
	Class(x) = argmax_c{P(y = c) * ∏P(x|y = c)}
And we can use logarithmic transformations to
make this calculation more computer-practical
(multiplying a bunch of probabilities on [0,1]
will always result in a very small number
which could easily underflow the float value):
	Class(x) = argmax_c{log(P(y = c)) + ΣP(x|y = c)}
Much better. That's our model!
*/
type NaiveBayes struct {
	// Words holds a map of words
	// to their corresponding Word
	// structure
	Words histogram `json:"words"`

	// Count holds the number of times
	// class i was seen as Count[i]
	Count []uint64 `json:"count"`

	// Probabilities holds the probability
	// that class Y is class i as
	// Probabilities[i] for
	Probabilities []float64 `json:"probabilities"`

	// DocumentCount holds the number of
	// documents that have been seen
	DocumentCount uint64 `json:"document_count"`

	// DictCount holds the size of the
	// NaiveBayes model's vocabulary
	DictCount uint64 `json:"vocabulary_size"`

	// sanitize is used by a model
	// to sanitize input of text
	sanitize transform.Transformer

	// stream holds the datastream
	stream <-chan base.TextDatapoint

	// Output is the io.Writer used for logging
	// and printing. Defaults to os.Stdout.
	Output io.Writer
}

// histogram allows conncurrency-friendly map access via its
// exported Get and Set methods
type histogram struct {
	sync.RWMutex
	words map[string]Word
}

// Get looks up a word from h's Word map, it should be used in
// place of a direct map lookup
// the only caveat here is that it will always return the 'success' boolean
func (h *histogram) Get(w string) (Word, bool) {
	h.RLock()
	result, ok := h.words[w]
	h.RUnlock()
	return result, ok
}

// Set sets word k's value to v in h's Word map
func (h *histogram) Set(k string, v Word) {
	h.Lock()
	h.words[k] = v
	h.Unlock()
}

// Word holds the structural
// information needed to calculate
// the probability of
type Word struct {
	// Count holds the number of times,
	// (i in Count[i] is the given class)
	Count []uint64

	// Seen holds the number of times
	// the world has been seen. This
	// is than same as
	//    foldl (+) 0 Count
	// in Haskell syntax, but is included
	// you wouldn't have to calculate
	// this every time you wanted to
	// recalc the probabilities (foldl
	// is the same as reduce, basically.)
	Seen uint64

	// DocsSeen is the same as Seen but
	// a word is only counted once even
	// if it's in a document multiple times
	DocsSeen uint64
}

// NewNaiveBayes returns a NaiveBayes model the
// given number of classes instantiated, ready
// to learn off the given data stream. The sanitization
// function is set to the given function. It must
// comply with the transform.RemoveFunc interface
func NewNaiveBayes(stream <-chan base.TextDatapoint, classes uint8, sanitize func(rune) bool) *NaiveBayes {
	return &NaiveBayes{
		Words:         histogram{sync.RWMutex{}, make(map[string]Word)},
		Count:         make([]uint64, classes),
		Probabilities: make([]float64, classes),

		sanitize: transform.RemoveFunc(sanitize),
		stream:   stream,

		Output: os.Stdout,
	}
}

// Predict takes in a document, predicts the
// class of the document based on the training
// data passed so far, and returns the class
// estimated for the document.
func (b *NaiveBayes) Predict(sentence string) uint8 {
	sums := make([]float64, len(b.Count))

	sentence, _, _ = transform.String(b.sanitize, sentence)
	words := strings.Split(strings.ToLower(sentence), " ")
	for _, word := range words {
		w, ok := b.Words.Get(word)
		if !ok {
			continue
		}

		for i := range sums {
			sums[i] += math.Log(float64(w.Count[i]+1) / float64(w.Seen+b.DictCount))
		}
	}

	for i := range sums {
		sums[i] += math.Log(b.Probabilities[i])
	}

	// find best class
	var maxI int
	for i := range sums {
		if sums[i] > sums[maxI] {
			maxI = i
		}
	}

	return uint8(maxI)
}

// Probability takes in a small document, returns the
// estimated class of the document based on the model
// as well as the probability that the model is part
// of that class
//
// NOTE: you should only use this for small documents
// because, as discussed in the docs for the model, the
// probability will often times underflow because you
// are multiplying together a bunch of probabilities
// which range on [0,1]. As such, the returned float
// could be NaN, and the predicted class could be
// 0 always.
//
// Basically, use Predict to be robust for larger
// documents. Use Probability only on relatively small
// (MAX of maybe a dozen words - basically just
// sentences and words) documents.
func (b *NaiveBayes) Probability(sentence string) (uint8, float64) {
	sums := make([]float64, len(b.Count))
	for i := range sums {
		sums[i] = 1
	}

	sentence, _, _ = transform.String(b.sanitize, sentence)
	words := strings.Split(strings.ToLower(sentence), " ")
	for _, word := range words {
		w, ok := b.Words.Get(word)
		if !ok {
			continue
		}

		for i := range sums {
			sums[i] *= float64(w.Count[i]+1) / float64(w.Seen+b.DictCount)
		}
	}

	for i := range sums {
		sums[i] *= b.Probabilities[i]
	}

	var denom float64
	var maxI int
	for i := range sums {
		if sums[i] > sums[maxI] {
			maxI = i
		}

		denom += sums[i]
	}

	return uint8(maxI), sums[maxI] / denom
}

// OnlineLearn lets the NaiveBayes model learn
// from the datastream, waiting for new data to
// come into the stream from a separate goroutine
func (b *NaiveBayes) OnlineLearn(errors chan<- error) {
	if errors == nil {
		errors = make(chan error)
	}
	if b.stream == nil {
		errors <- fmt.Errorf("ERROR: attempting to learn with nil data stream!\n")
		close(errors)
		return
	}

	fmt.Fprintf(b.Output, "Training:\n\tModel: Multinomial Naïve Bayes\n\tClasses: %v\n", len(b.Count))

	var point base.TextDatapoint
	var more bool

	for {
		point, more = <-b.stream

		if more {
			// sanitize and break up document
			sanitized, _, _ := transform.String(b.sanitize, point.X)
			sanitized = strings.ToLower(sanitized)

			words := strings.Split(sanitized, " ")

			C := int(point.Y)

			if C > len(b.Count)-1 {
				errors <- fmt.Errorf("ERROR: given document class is greater than the number of classes in the model!\n")
				continue
			}

			// update global class probabilities
			b.Count[C]++
			b.DocumentCount++
			for i := range b.Probabilities {
				b.Probabilities[i] = float64(b.Count[i]) / float64(b.DocumentCount)
			}

			// store words seen in document (to add to DocsSeen)
			seenCount := make(map[string]int)

			// update probabilities for words
			for _, word := range words {
				if len(word) < 3 {
					continue
				}

				w, ok := b.Words.Get(word)

				if !ok {
					w = Word{
						Count: make([]uint64, len(b.Count)),
						Seen:  uint64(0),
					}

					b.DictCount++
				}

				w.Count[C]++
				w.Seen++

				b.Words.Set(word, w)

				seenCount[word] = 1
			}

			// add to DocsSeen
			for term := range seenCount {
				tmp, _ := b.Words.Get(term)
				tmp.DocsSeen++
				b.Words.Set(term, tmp)
			}
		} else {
			fmt.Fprintf(b.Output, "Training Completed.\n%v\n\n", b)
			close(errors)
			return
		}
	}
}

// UpdateStream updates the NaiveBayes model's
// text datastream
func (b *NaiveBayes) UpdateStream(stream chan base.TextDatapoint) {
	b.stream = stream
}

// UpdateSanitize updates the NaiveBayes model's
// text sanitization transformation function
func (b *NaiveBayes) UpdateSanitize(sanitize func(rune) bool) {
	b.sanitize = transform.RemoveFunc(sanitize)
}

// String implements the fmt interface for clean printing. Here
// we're using it to print the model as the equation h(θ)=...
// where h is the perceptron hypothesis model.
func (b *NaiveBayes) String() string {
	return fmt.Sprintf("h(θ) = argmax_c{log(P(y = c)) + Σlog(P(x|y = c))}\n\tClasses: %v\n\tDocuments evaluated in model: %v\n\tWords evaluated in model: %v\n", len(b.Count), int(b.DocumentCount), int(b.DictCount))
}

// PersistToFile takes in an absolute filepath and saves the
// parameter vector θ to the file, which can be restored later.
// The function will take paths from the current directory, but
// functions
//
// The data is stored as JSON because it's one of the most
// efficient storage method (you only need one comma extra
// per feature + two brackets, total!) And it's extendable.
func (b *NaiveBayes) PersistToFile(path string) error {
	if path == "" {
		return fmt.Errorf("ERROR: you just tried to persist your model to a file with no path!! That's a no-no. Try it with a valid filepath")
	}

	bytes, err := json.Marshal(b)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(path, bytes, os.ModePerm)
	if err != nil {
		return err
	}

	return nil
}

// Restore takes the bytes of a NaiveBayes model and
// restores a model to it. This would be useful if
// training a model and saving it into a project
// using go-bindata (look it up) so you don't have
// to persist a large file and deal with paths on
// a production system. This option is included
// in text models vs. others because the text models
// usually have much larger storage requirements
func (b *NaiveBayes) Restore(bytes []byte) error {
	err := json.Unmarshal(bytes, &b)
	if err != nil {
		return err
	}

	return nil
}

// RestoreFromFile takes in a path to a parameter vector theta
// and assigns the model it's operating on's parameter vector
// to that.
//
// The path must ba an absolute path or a path from the current
// directory
//
// This would be useful in persisting data between running
// a model on data.
func (b *NaiveBayes) RestoreFromFile(path string) error {
	if path == "" {
		return fmt.Errorf("ERROR: you just tried to restore your model from a file with no path! That's a no-no. Try it with a valid filepath")
	}

	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	err = json.Unmarshal(bytes, &b)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package service

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
)

type RemoteCatalogClient struct {
	serverEndpoint *url.URL
}

func serviceFromResponse(res *http.Response, apiLocation string) (*Service, error) {
	decoder := json.NewDecoder(res.Body)
	defer res.Body.Close()

	var s *Service
	err := decoder.Decode(&s)
	if err != nil {
		return nil, err
	}
	svc := s.unLdify(apiLocation)
	return &svc, nil
}

func NewRemoteCatalogClient(serverEndpoint string) *RemoteCatalogClient {
	// Check if serverEndpoint is a correct URL
	endpointUrl, err := url.Parse(serverEndpoint)
	if err != nil {
		return &RemoteCatalogClient{}
	}

	return &RemoteCatalogClient{
		serverEndpoint: endpointUrl,
	}
}

func (self *RemoteCatalogClient) Get(id string) (*Service, error) {
	res, err := http.Get(fmt.Sprintf("%v/%v", self.serverEndpoint, id))
	if err != nil {
		return nil, err
	}

	if res.StatusCode == http.StatusNotFound {
		return nil, ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%v", res.StatusCode)
	}
	return serviceFromResponse(res, self.serverEndpoint.Path)
}

func (self *RemoteCatalogClient) Add(s *Service) error {
	b, _ := json.Marshal(s)
	_, err := http.Post(self.serverEndpoint.String()+"/", "application/ld+json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	return nil
}

func (self *RemoteCatalogClient) Update(id string, s *Service) error {
	b, _ := json.Marshal(s)
	req, err := http.NewRequest("PUT", fmt.Sprintf("%v/%v", self.serverEndpoint, id), bytes.NewReader(b))
	if err != nil {
		return err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	if res.StatusCode == http.StatusNotFound {
		return ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%v", res.StatusCode)
	}
	return nil
}

func (self *RemoteCatalogClient) Delete(id string) error {
	req, err := http.NewRequest("DELETE", fmt.Sprintf("%v/%v", self.serverEndpoint, id), bytes.NewReader([]byte{}))
	if err != nil {
		return err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	if res.StatusCode == http.StatusNotFound {
		return ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%v", res.StatusCode)
	}

	return nil
}

func (self *RemoteCatalogClient) GetServices(page, perPage int) ([]Service, int, error) {
	res, err := http.Get(
		fmt.Sprintf("%s?%s=%s&%s=%s",
			self.serverEndpoint, GetParamPage, page, GetParamPerPage, perPage))
	if err != nil {
		return nil, 0, err
	}

	decoder := json.NewDecoder(res.Body)
	defer res.Body.Close()

	var coll Collection
	err = decoder.Decode(&coll)
	if err != nil {
		return nil, 0, err
	}

	svcs := make([]Service, 0, len(coll.Services))
	for _, v := range coll.Services {
		svcs = append(svcs, v.unLdify(self.serverEndpoint.Path))
	}

	return svcs, len(svcs), nil
}

// TODO
func (self *RemoteCatalogClient) FindService(path, op, value string) (*Service, error) {
	return nil, nil
}

// TODO
func (self *RemoteCatalogClient) FindServices(path, op, value string, page, perPage int) ([]Service, int, error) {
	return nil, 0, nil
}
<commit_msg>FindService implementation for #5<commit_after>package service

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
)

type RemoteCatalogClient struct {
	serverEndpoint *url.URL
}

func serviceFromResponse(res *http.Response, apiLocation string) (*Service, error) {
	decoder := json.NewDecoder(res.Body)
	defer res.Body.Close()

	var s *Service
	err := decoder.Decode(&s)
	if err != nil {
		return nil, err
	}
	svc := s.unLdify(apiLocation)
	return &svc, nil
}

func NewRemoteCatalogClient(serverEndpoint string) *RemoteCatalogClient {
	// Check if serverEndpoint is a correct URL
	endpointUrl, err := url.Parse(serverEndpoint)
	if err != nil {
		return &RemoteCatalogClient{}
	}

	return &RemoteCatalogClient{
		serverEndpoint: endpointUrl,
	}
}

func (self *RemoteCatalogClient) Get(id string) (*Service, error) {
	res, err := http.Get(fmt.Sprintf("%v/%v", self.serverEndpoint, id))
	if err != nil {
		return nil, err
	}

	if res.StatusCode == http.StatusNotFound {
		return nil, ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%v", res.StatusCode)
	}
	return serviceFromResponse(res, self.serverEndpoint.Path)
}

func (self *RemoteCatalogClient) Add(s *Service) error {
	b, _ := json.Marshal(s)
	_, err := http.Post(self.serverEndpoint.String()+"/", "application/ld+json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	return nil
}

func (self *RemoteCatalogClient) Update(id string, s *Service) error {
	b, _ := json.Marshal(s)
	req, err := http.NewRequest("PUT", fmt.Sprintf("%v/%v", self.serverEndpoint, id), bytes.NewReader(b))
	if err != nil {
		return err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	if res.StatusCode == http.StatusNotFound {
		return ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%v", res.StatusCode)
	}
	return nil
}

func (self *RemoteCatalogClient) Delete(id string) error {
	req, err := http.NewRequest("DELETE", fmt.Sprintf("%v/%v", self.serverEndpoint, id), bytes.NewReader([]byte{}))
	if err != nil {
		return err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	if res.StatusCode == http.StatusNotFound {
		return ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%v", res.StatusCode)
	}

	return nil
}

func (self *RemoteCatalogClient) GetServices(page, perPage int) ([]Service, int, error) {
	res, err := http.Get(
		fmt.Sprintf("%s?%s=%s&%s=%s",
			self.serverEndpoint, GetParamPage, page, GetParamPerPage, perPage))
	if err != nil {
		return nil, 0, err
	}

	decoder := json.NewDecoder(res.Body)
	defer res.Body.Close()

	var coll Collection
	err = decoder.Decode(&coll)
	if err != nil {
		return nil, 0, err
	}

	svcs := make([]Service, 0, len(coll.Services))
	for _, v := range coll.Services {
		svcs = append(svcs, v.unLdify(self.serverEndpoint.Path))
	}

	return svcs, len(svcs), nil
}

func (self *RemoteCatalogClient) FindService(path, op, value string) (*Service, error) {
	res, err := http.Get(fmt.Sprintf("%v/%v/%v/%v", self.serverEndpoint, path, op, value))
	if err != nil {
		return nil, err
	}

	if res.StatusCode == http.StatusNotFound {
		return nil, ErrorNotFound
	} else if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%v", res.StatusCode)
	}
	return serviceFromResponse(res, self.serverEndpoint.Path)
}

// TODO
func (self *RemoteCatalogClient) FindServices(path, op, value string, page, perPage int) ([]Service, int, error) {
	return nil, 0, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by
// license that can be found in the LICENSE file.

// Package daemon windows version
package daemon

import (
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"syscall"
	"time"
	"unicode/utf16"
	"unsafe"

	"golang.org/x/sys/windows/registry"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/mgr"
)

// windowsRecord - standard record (struct) for windows version of daemon package
type windowsRecord struct {
	name         string
	description  string
	dependencies []string
}

func newDaemon(name, description string, dependencies []string) (Daemon, error) {

	return &windowsRecord{name, description, dependencies}, nil
}

// Install the service
func (windows *windowsRecord) Install(args ...string) (string, error) {
	installAction := "Install " + windows.description + ":"

	execp, err := execPath()

	if err != nil {
		return installAction + failed, err
	}

	m, err := mgr.Connect()
	if err != nil {
		return installAction + failed, err
	}
	defer m.Disconnect()

	s, err := m.OpenService(windows.name)
	if err == nil {
		s.Close()
		return installAction + failed, err
	}

	s, err = m.CreateService(windows.name, execp, mgr.Config{
		DisplayName:  windows.name,
		Description:  windows.description,
		StartType:    mgr.StartAutomatic,
		Dependencies: windows.dependencies,
	}, args...)
	if err != nil {
		return installAction + failed, err
	}
	defer s.Close()

	return installAction + " completed.", nil
}

// Remove the service
func (windows *windowsRecord) Remove() (string, error) {
	removeAction := "Removing " + windows.description + ":"

	m, err := mgr.Connect()
	if err != nil {
		return removeAction + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return removeAction + failed, getWindowsError(err)
	}
	defer s.Close()
	err = s.Delete()
	if err != nil {
		return removeAction + failed, getWindowsError(err)
	}

	return removeAction + " completed.", nil
}

// Start the service
func (windows *windowsRecord) Start() (string, error) {
	startAction := "Starting " + windows.description + ":"

	m, err := mgr.Connect()
	if err != nil {
		return startAction + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return startAction + failed, getWindowsError(err)
	}
	defer s.Close()
	if err = s.Start(); err != nil {
		return startAction + failed, getWindowsError(err)
	}

	return startAction + " completed.", nil
}

// Stop the service
func (windows *windowsRecord) Stop() (string, error) {
	stopAction := "Stopping " + windows.description + ":"

	m, err := mgr.Connect()
	if err != nil {
		return stopAction + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return stopAction + failed, getWindowsError(err)
	}
	defer s.Close()
	if err := stopAndWait(s); err != nil {
		return stopAction + failed, getWindowsError(err)
	}

	return stopAction + " completed.", nil
}

func stopAndWait(s *mgr.Service) error {
	// First stop the service. Then wait for the service to
	// actually stop before starting it.
	status, err := s.Control(svc.Stop)
	if err != nil {
		return err
	}

	timeDuration := time.Millisecond * 50

	timeout := time.After(getStopTimeout() + (timeDuration * 2))
	tick := time.NewTicker(timeDuration)
	defer tick.Stop()

	for status.State != svc.Stopped {
		select {
		case <-tick.C:
			status, err = s.Query()
			if err != nil {
				return err
			}
		case <-timeout:
			break
		}
	}
	return nil
}

func getStopTimeout() time.Duration {
	// For default and paths see https://support.microsoft.com/en-us/kb/146092
	defaultTimeout := time.Millisecond * 20000
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\CurrentControlSet\Control`, registry.READ)
	if err != nil {
		return defaultTimeout
	}
	sv, _, err := key.GetStringValue("WaitToKillServiceTimeout")
	if err != nil {
		return defaultTimeout
	}
	v, err := strconv.Atoi(sv)
	if err != nil {
		return defaultTimeout
	}
	return time.Millisecond * time.Duration(v)
}

// Status - Get service status
func (windows *windowsRecord) Status() (string, error) {
	m, err := mgr.Connect()
	if err != nil {
		return "Getting status:" + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return "Getting status:" + failed, getWindowsError(err)
	}
	defer s.Close()
	status, err := s.Query()
	if err != nil {
		return "Getting status:" + failed, getWindowsError(err)
	}

	return "Status: " + getWindowsServiceStateFromUint32(status.State), nil
}

// Get executable path
func execPath() (string, error) {
	var n uint32
	b := make([]uint16, syscall.MAX_PATH)
	size := uint32(len(b))

	r0, _, e1 := syscall.MustLoadDLL(
		"kernel32.dll",
	).MustFindProc(
		"GetModuleFileNameW",
	).Call(0, uintptr(unsafe.Pointer(&b[0])), uintptr(size))
	n = uint32(r0)
	if n == 0 {
		return "", e1
	}
	return string(utf16.Decode(b[0:n])), nil
}

// Get windows error
func getWindowsError(inputError error) error {
	if exiterr, ok := inputError.(*exec.ExitError); ok {
		if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
			if sysErr, ok := WinErrCode[status.ExitStatus()]; ok {
				return errors.New(fmt.Sprintf("\n %s: %s \n %s", sysErr.Title, sysErr.Description, sysErr.Action))
			}
		}
	}

	return inputError
}

// Get windows service state
func getWindowsServiceStateFromUint32(state svc.State) string {
	switch state {
	case svc.Stopped:
		return "SERVICE_STOPPED"
	case svc.StartPending:
		return "SERVICE_START_PENDING"
	case svc.StopPending:
		return "SERVICE_STOP_PENDING"
	case svc.Running:
		return "SERVICE_RUNNING"
	case svc.ContinuePending:
		return "SERVICE_CONTINUE_PENDING"
	case svc.PausePending:
		return "SERVICE_PAUSE_PENDING"
	case svc.Paused:
		return "SERVICE_PAUSED"
	}
	return "SERVICE_UNKNOWN"
}

type serviceHandler struct {
	executable Executable
}

func (sh *serviceHandler) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown | svc.AcceptPauseAndContinue
	changes <- svc.Status{State: svc.StartPending}

	fasttick := time.Tick(500 * time.Millisecond)
	slowtick := time.Tick(2 * time.Second)
	tick := fasttick

	sh.executable.Start()
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}

loop:
	for {
		select {
		case <-tick:
			break
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
				// Testing deadlock from https://code.google.com/p/winsvc/issues/detail?id=4
				time.Sleep(100 * time.Millisecond)
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				changes <- svc.Status{State: svc.StopPending}
				sh.executable.Stop()
				break loop
			case svc.Pause:
				changes <- svc.Status{State: svc.Paused, Accepts: cmdsAccepted}
				tick = slowtick
			case svc.Continue:
				changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
				tick = fasttick
			default:
				continue loop
			}
		}
	}
	return
}

func (windows *windowsRecord) Run(e Executable) (string, error) {
	runAction := "Running " + windows.description + ":"

	interactive, err := svc.IsAnInteractiveSession()
	if err != nil {
		return runAction + failed, getWindowsError(err)
	}
	if interactive {
		e.Run()
	}

	err = svc.Run(windows.name, &serviceHandler{
		executable: e,
	})
	if err != nil {
		return runAction + failed, getWindowsError(err)
	}

	return runAction + " completed.", nil
}
<commit_msg>fix if condition bug and add comments<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by
// license that can be found in the LICENSE file.

// Package daemon windows version
package daemon

import (
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"syscall"
	"time"
	"unicode/utf16"
	"unsafe"

	"golang.org/x/sys/windows/registry"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/mgr"
)

// windowsRecord - standard record (struct) for windows version of daemon package
type windowsRecord struct {
	name         string
	description  string
	dependencies []string
}

func newDaemon(name, description string, dependencies []string) (Daemon, error) {

	return &windowsRecord{name, description, dependencies}, nil
}

// Install the service
func (windows *windowsRecord) Install(args ...string) (string, error) {
	installAction := "Install " + windows.description + ":"

	execp, err := execPath()

	if err != nil {
		return installAction + failed, err
	}

	m, err := mgr.Connect()
	if err != nil {
		return installAction + failed, err
	}
	defer m.Disconnect()

	s, err := m.OpenService(windows.name)
	if err == nil {
		s.Close()
		return installAction + failed, err
	}

	s, err = m.CreateService(windows.name, execp, mgr.Config{
		DisplayName:  windows.name,
		Description:  windows.description,
		StartType:    mgr.StartAutomatic,
		Dependencies: windows.dependencies,
	}, args...)
	if err != nil {
		return installAction + failed, err
	}
	defer s.Close()

	return installAction + " completed.", nil
}

// Remove the service
func (windows *windowsRecord) Remove() (string, error) {
	removeAction := "Removing " + windows.description + ":"

	m, err := mgr.Connect()
	if err != nil {
		return removeAction + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return removeAction + failed, getWindowsError(err)
	}
	defer s.Close()
	err = s.Delete()
	if err != nil {
		return removeAction + failed, getWindowsError(err)
	}

	return removeAction + " completed.", nil
}

// Start the service
func (windows *windowsRecord) Start() (string, error) {
	startAction := "Starting " + windows.description + ":"

	m, err := mgr.Connect()
	if err != nil {
		return startAction + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return startAction + failed, getWindowsError(err)
	}
	defer s.Close()
	if err = s.Start(); err != nil {
		return startAction + failed, getWindowsError(err)
	}

	return startAction + " completed.", nil
}

// Stop the service
func (windows *windowsRecord) Stop() (string, error) {
	stopAction := "Stopping " + windows.description + ":"

	m, err := mgr.Connect()
	if err != nil {
		return stopAction + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return stopAction + failed, getWindowsError(err)
	}
	defer s.Close()
	if err := stopAndWait(s); err != nil {
		return stopAction + failed, getWindowsError(err)
	}

	return stopAction + " completed.", nil
}

func stopAndWait(s *mgr.Service) error {
	// First stop the service. Then wait for the service to
	// actually stop before starting it.
	status, err := s.Control(svc.Stop)
	if err != nil {
		return err
	}

	timeDuration := time.Millisecond * 50

	timeout := time.After(getStopTimeout() + (timeDuration * 2))
	tick := time.NewTicker(timeDuration)
	defer tick.Stop()

	for status.State != svc.Stopped {
		select {
		case <-tick.C:
			status, err = s.Query()
			if err != nil {
				return err
			}
		case <-timeout:
			break
		}
	}
	return nil
}

func getStopTimeout() time.Duration {
	// For default and paths see https://support.microsoft.com/en-us/kb/146092
	defaultTimeout := time.Millisecond * 20000
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\CurrentControlSet\Control`, registry.READ)
	if err != nil {
		return defaultTimeout
	}
	sv, _, err := key.GetStringValue("WaitToKillServiceTimeout")
	if err != nil {
		return defaultTimeout
	}
	v, err := strconv.Atoi(sv)
	if err != nil {
		return defaultTimeout
	}
	return time.Millisecond * time.Duration(v)
}

// Status - Get service status
func (windows *windowsRecord) Status() (string, error) {
	m, err := mgr.Connect()
	if err != nil {
		return "Getting status:" + failed, getWindowsError(err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(windows.name)
	if err != nil {
		return "Getting status:" + failed, getWindowsError(err)
	}
	defer s.Close()
	status, err := s.Query()
	if err != nil {
		return "Getting status:" + failed, getWindowsError(err)
	}

	return "Status: " + getWindowsServiceStateFromUint32(status.State), nil
}

// Get executable path
func execPath() (string, error) {
	var n uint32
	b := make([]uint16, syscall.MAX_PATH)
	size := uint32(len(b))

	r0, _, e1 := syscall.MustLoadDLL(
		"kernel32.dll",
	).MustFindProc(
		"GetModuleFileNameW",
	).Call(0, uintptr(unsafe.Pointer(&b[0])), uintptr(size))
	n = uint32(r0)
	if n == 0 {
		return "", e1
	}
	return string(utf16.Decode(b[0:n])), nil
}

// Get windows error
func getWindowsError(inputError error) error {
	if exiterr, ok := inputError.(*exec.ExitError); ok {
		if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
			if sysErr, ok := WinErrCode[status.ExitStatus()]; ok {
				return errors.New(fmt.Sprintf("\n %s: %s \n %s", sysErr.Title, sysErr.Description, sysErr.Action))
			}
		}
	}

	return inputError
}

// Get windows service state
func getWindowsServiceStateFromUint32(state svc.State) string {
	switch state {
	case svc.Stopped:
		return "SERVICE_STOPPED"
	case svc.StartPending:
		return "SERVICE_START_PENDING"
	case svc.StopPending:
		return "SERVICE_STOP_PENDING"
	case svc.Running:
		return "SERVICE_RUNNING"
	case svc.ContinuePending:
		return "SERVICE_CONTINUE_PENDING"
	case svc.PausePending:
		return "SERVICE_PAUSE_PENDING"
	case svc.Paused:
		return "SERVICE_PAUSED"
	}
	return "SERVICE_UNKNOWN"
}

type serviceHandler struct {
	executable Executable
}

func (sh *serviceHandler) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown | svc.AcceptPauseAndContinue
	changes <- svc.Status{State: svc.StartPending}

	fasttick := time.Tick(500 * time.Millisecond)
	slowtick := time.Tick(2 * time.Second)
	tick := fasttick

	sh.executable.Start()
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}

loop:
	for {
		select {
		case <-tick:
			break
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
				// Testing deadlock from https://code.google.com/p/winsvc/issues/detail?id=4
				time.Sleep(100 * time.Millisecond)
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				changes <- svc.Status{State: svc.StopPending}
				sh.executable.Stop()
				break loop
			case svc.Pause:
				changes <- svc.Status{State: svc.Paused, Accepts: cmdsAccepted}
				tick = slowtick
			case svc.Continue:
				changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
				tick = fasttick
			default:
				continue loop
			}
		}
	}
	return
}

func (windows *windowsRecord) Run(e Executable) (string, error) {
	runAction := "Running " + windows.description + ":"

	interactive, err := svc.IsAnInteractiveSession()
	if err != nil {
		return runAction + failed, getWindowsError(err)
	}
	if !interactive {
		// service called from windows service manager
		// use API provided by golang.org/x/sys/windows
		err = svc.Run(windows.name, &serviceHandler{
			executable: e,
		})
		if err != nil {
			return runAction + failed, getWindowsError(err)
		}
	} else {
		// otherwise, service should be called from terminal session
		e.Run()
	}

	return runAction + " completed.", nil
}
<|endoftext|>


<commit_before>package tile

import (
	"math"
	"time"
	"fmt"
	"strconv"
	"sync"
)

func Tiler(g Grid, maxZoom int, targetFolder string){
	defer TimeTrack(time.Now(), fmt.Sprint("Tiler with zoom ", maxZoom))

	tileSize := 256
	world := Bounds{
		Point{MercatorPole*-1,MercatorPole*-1},
		Point{MercatorPole,MercatorPole}}

	totalNumTiles:=4*(math.Pow(4,float64(maxZoom))-1)/3

	var wg sync.WaitGroup
	wg.Add(int(totalNumTiles))

	fmt.Println(totalNumTiles)

	for z:=1; z<=maxZoom; z++ {
		sizef := math.Pow(2,float64(z))
		for x:=0; x<int(sizef); x++ {
			xf := float64(x)
			for y:=0; y<int(sizef); y++ {
				width := world.Width() / sizef
				yf := float64(y)

				//make bounds
				b := Bounds{
					Point{width * xf, MercatorPole - width * (yf + 1)},
					Point{width * (xf + 1), MercatorPole - width * yf}}

				fileName := fmt.Sprintf(targetFolder + "\\" + strconv.Itoa(int(z)) + "_" + strconv.Itoa(int(x)) + "_" + strconv.Itoa(int(y)) + ".png")

				go (func(b Bounds, fileName string) {

					//make tile
					gTile := Resample(g, b, tileSize, tileSize)

					//make image
					img := DrawGrid(&gTile)

					//save image to disc
					SaveImage(img, fileName)

					wg.Done()
				})(b,fileName)
			}
		}
	}

	wg.Wait()

}<commit_msg>Example on creating tiles based on grib2json file<commit_after>package tile

import (
	"math"
	"time"
	"fmt"
	"strconv"
	"sync"
)

func Tiler(g Grid, maxZoom int, targetFolder string){
	defer TimeTrack(time.Now(), fmt.Sprint("Tiler with zoom ", maxZoom))

	tileSize := 256
	world := Bounds{
		Point{MercatorPole*-1,MercatorPole*-1},
		Point{MercatorPole,MercatorPole}}

	var wg sync.WaitGroup

	for z:=1; z<=maxZoom; z++ {
		sizef := math.Pow(2,float64(z))
		for x:=0; x<int(sizef); x++ {
			xf := float64(x)
			for y:=0; y<int(sizef); y++ {
				wg.Add(1)

				width := world.Width() / sizef
				yf := float64(y)

				//make bounds
				b := Bounds{
					Point{MercatorPole*-1 +  width * xf, MercatorPole - width * (yf + 1)},
					Point{MercatorPole*-1 +  width * (xf + 1), MercatorPole - width * yf}}

				fileName := fmt.Sprintf(targetFolder + "\\" + strconv.Itoa(int(z)) + "_" + strconv.Itoa(int(x)) + "_" + strconv.Itoa(int(y)) + ".png")

				go (func(b Bounds, fileName string) {

					//make tile
					gTile := Resample(g, b, tileSize, tileSize)

					//make image
					img := DrawGrid(&gTile)

					//save image to disc
					SaveImage(img, fileName)

					wg.Done()
				})(b,fileName)
			}
		}
	}

	wg.Wait()

}<|endoftext|>


<commit_before>package kcp

import (
	"container/heap"
	"runtime"
	"sync"
	"time"
)

// SystemTimedSched is the library level timed-scheduler
var SystemTimedSched *TimedSched = NewTimedSched(runtime.NumCPU())

type timedFunc struct {
	execute func()
	ts      time.Time
}

// a heap for sorted timed function
type timedFuncHeap []timedFunc

func (h timedFuncHeap) Len() int            { return len(h) }
func (h timedFuncHeap) Less(i, j int) bool  { return h[i].ts.Before(h[j].ts) }
func (h timedFuncHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *timedFuncHeap) Push(x interface{}) { *h = append(*h, x.(timedFunc)) }
func (h *timedFuncHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	old[n-1].execute = nil // avoid memory leak
	*h = old[0 : n-1]
	return x
}

// TimedSched represents the control struct for timed parallel scheduler
type TimedSched struct {
	// prepending tasks
	prependTasks    []timedFunc
	prependLock     sync.Mutex
	chPrependNotify chan struct{}

	// tasks will be distributed through chTask
	chTask chan timedFunc

	dieOnce sync.Once
	die     chan struct{}
}

// NewTimedSched creates a parallel-scheduler with given parallelization
func NewTimedSched(parallel int) *TimedSched {
	ts := new(TimedSched)
	ts.chTask = make(chan timedFunc)
	ts.die = make(chan struct{})
	ts.chPrependNotify = make(chan struct{}, 1)

	for i := 0; i < parallel; i++ {
		go ts.sched()
	}
	go ts.prepend()
	return ts
}

func (ts *TimedSched) sched() {
	var tasks timedFuncHeap
	timer := time.NewTimer(0)
	for {
		select {
		case task := <-ts.chTask:
			now := time.Now()
			if now.After(task.ts) {
				// already delayed! execute immediately
				task.execute()
			} else {
				heap.Push(&tasks, task)
				// reset timer to trigger based on the top element
				timer.Reset(tasks[0].ts.Sub(now))
			}
		case <-timer.C:
			for tasks.Len() > 0 {
				now := time.Now()
				if now.After(tasks[0].ts) {
					heap.Pop(&tasks).(timedFunc).execute()
				} else {
					timer.Reset(tasks[0].ts.Sub(now))
					break
				}
			}
		case <-ts.die:
			return
		}
	}
}

func (ts *TimedSched) prepend() {
	var tasks []timedFunc
	for {
		select {
		case <-ts.chPrependNotify:
			ts.prependLock.Lock()
			// keep cap to reuse slice
			if cap(tasks) < cap(ts.prependTasks) {
				tasks = make([]timedFunc, 0, cap(ts.prependTasks))
			}
			tasks = tasks[:len(ts.prependTasks)]
			copy(tasks, ts.prependTasks)
			for k := range ts.prependTasks {
				ts.prependTasks[k].execute = nil // avoid memory leak
			}
			ts.prependTasks = ts.prependTasks[:0]
			ts.prependLock.Unlock()

			for k := range tasks {
				select {
				case ts.chTask <- tasks[k]:
					tasks[k].execute = nil // avoid memory leak
				case <-ts.die:
					return
				}
			}
			tasks = tasks[:0]
		case <-ts.die:
			return
		}
	}
}

// Put a function 'f' awaiting to be executed at 'deadline'
func (ts *TimedSched) Put(f func(), deadline time.Time) {
	ts.prependLock.Lock()
	ts.prependTasks = append(ts.prependTasks, timedFunc{f, deadline})
	ts.prependLock.Unlock()

	select {
	case ts.chPrependNotify <- struct{}{}:
	default:
	}
}

// Close terminates this scheduler
func (ts *TimedSched) Close() { ts.dieOnce.Do(func() { close(ts.die) }) }
<commit_msg>optimize on timedsched<commit_after>package kcp

import (
	"container/heap"
	"runtime"
	"sync"
	"time"
)

// SystemTimedSched is the library level timed-scheduler
var SystemTimedSched *TimedSched = NewTimedSched(runtime.NumCPU())

type timedFunc struct {
	execute func()
	ts      time.Time
}

// a heap for sorted timed function
type timedFuncHeap []timedFunc

func (h timedFuncHeap) Len() int            { return len(h) }
func (h timedFuncHeap) Less(i, j int) bool  { return h[i].ts.Before(h[j].ts) }
func (h timedFuncHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *timedFuncHeap) Push(x interface{}) { *h = append(*h, x.(timedFunc)) }
func (h *timedFuncHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	old[n-1].execute = nil // avoid memory leak
	*h = old[0 : n-1]
	return x
}

// TimedSched represents the control struct for timed parallel scheduler
type TimedSched struct {
	// prepending tasks
	prependTasks    []timedFunc
	prependLock     sync.Mutex
	chPrependNotify chan struct{}

	// tasks will be distributed through chTask
	chTask chan timedFunc

	dieOnce sync.Once
	die     chan struct{}
}

// NewTimedSched creates a parallel-scheduler with given parallelization
func NewTimedSched(parallel int) *TimedSched {
	ts := new(TimedSched)
	ts.chTask = make(chan timedFunc)
	ts.die = make(chan struct{})
	ts.chPrependNotify = make(chan struct{}, 1)

	for i := 0; i < parallel; i++ {
		go ts.sched()
	}
	go ts.prepend()
	return ts
}

func (ts *TimedSched) sched() {
	var tasks timedFuncHeap
	timer := time.NewTimer(0)
	for {
		select {
		case task := <-ts.chTask:
			now := time.Now()
			if now.After(task.ts) {
				// already delayed! execute immediately
				task.execute()
			} else {
				heap.Push(&tasks, task)
				// reset timer to trigger based on the top element
				timer.Reset(tasks[0].ts.Sub(now))
			}
		case now := <-timer.C:
			for tasks.Len() > 0 {
				if now.After(tasks[0].ts) {
					heap.Pop(&tasks).(timedFunc).execute()
				} else {
					timer.Reset(tasks[0].ts.Sub(now))
					break
				}
			}
		case <-ts.die:
			return
		}
	}
}

func (ts *TimedSched) prepend() {
	var tasks []timedFunc
	for {
		select {
		case <-ts.chPrependNotify:
			ts.prependLock.Lock()
			// keep cap to reuse slice
			if cap(tasks) < cap(ts.prependTasks) {
				tasks = make([]timedFunc, 0, cap(ts.prependTasks))
			}
			tasks = tasks[:len(ts.prependTasks)]
			copy(tasks, ts.prependTasks)
			for k := range ts.prependTasks {
				ts.prependTasks[k].execute = nil // avoid memory leak
			}
			ts.prependTasks = ts.prependTasks[:0]
			ts.prependLock.Unlock()

			for k := range tasks {
				select {
				case ts.chTask <- tasks[k]:
					tasks[k].execute = nil // avoid memory leak
				case <-ts.die:
					return
				}
			}
			tasks = tasks[:0]
		case <-ts.die:
			return
		}
	}
}

// Put a function 'f' awaiting to be executed at 'deadline'
func (ts *TimedSched) Put(f func(), deadline time.Time) {
	ts.prependLock.Lock()
	ts.prependTasks = append(ts.prependTasks, timedFunc{f, deadline})
	ts.prependLock.Unlock()

	select {
	case ts.chPrependNotify <- struct{}{}:
	default:
	}
}

// Close terminates this scheduler
func (ts *TimedSched) Close() { ts.dieOnce.Do(func() { close(ts.die) }) }
<|endoftext|>