<commit_before>package main

import (
	"errors"
	"github.com/phzfi/RIC/server/logging"
	"github.com/phzfi/RIC/server/ops"
	"github.com/valyala/fasthttp"
	"strings"
)

func ParseURI(uri *fasthttp.URI, source ops.ImageSource, marker ops.Watermarker) (operations []ops.Operation, format string, err, invalid error) {
	filename := string(uri.Path())

	w, h, cropx, cropy, mode, format, url, invalid := getParams(uri.QueryArgs())
	if invalid != nil {
		return
	}

	if url != "" {
		source.AddRoot(url)
	}

	ow, oh, err := source.ImageSize(filename)
	if err != nil {
		return
	}

	operations = []ops.Operation{source.LoadImageOp(filename)}

	adjustWidth := func() {
		w = roundedIntegerDivision(h*ow, oh)
	}

	adjustHeight := func() {
		h = roundedIntegerDivision(w*oh, ow)
	}

	adjustSize := func() {
		if h == 0 && w != 0 {
			adjustHeight()
		} else if h != 0 && w == 0 {
			adjustWidth()
		} else if w == 0 && h == 0 {
			w, h = ow, oh
		}
	}

	denyUpscale := func() {
		h0 := h
		w0 := w
		if w > ow {
			h = roundedIntegerDivision(ow*h0, w0)
			w = ow
		}
		if h > oh || h > h0 {
			w = roundedIntegerDivision(oh*w0, h0)
			h = oh
		}
	}

	resize := func() {
		denyUpscale()
		adjustSize()
		operations = append(operations, ops.Resize{w, h})
	}

	liquid := func() {
		denyUpscale()
		adjustSize()
		operations = append(operations, ops.LiquidRescale{w, h})
	}

	crop := func() {
		if w == 0 {
			w = ow
		}
		if h == 0 {
			h = oh
		}
		operations = append(operations, ops.Crop{w, h, cropx, cropy})
	}

	cropmid := func() {
		if w == 0 || w > ow {
			w = ow
		}
		if h == 0 || h > oh {
			h = oh
		}
		midW := roundedIntegerDivision(ow, 2)
		midH := roundedIntegerDivision(oh, 2)
		cropx := midW - roundedIntegerDivision(w, 2)
		cropy := midH - roundedIntegerDivision(h, 2)
		operations = append(operations, ops.Crop{w, h, cropx, cropy})
	}

	fit := func() {
		if w > ow {
			w = ow
		}
		if h > oh {
			h = oh
		}
		if w != 0 && h != 0 {
			if ow*h > w*oh {
				adjustHeight()
			} else {
				adjustWidth()
			}
			operations = append(operations, ops.Resize{w, h})
		} else {
			resize()
		}
	}

	watermark := func() {
		heightOK := h > marker.MinHeight && h < marker.MaxHeight
		widthOK := w > marker.MinWidth && w < marker.MaxWidth
		if marker.AddMark && heightOK && widthOK {
			logging.Debug("Adding watermarkOp")
			operations = append(operations, ops.WatermarkOp(marker.WatermarkImage, marker.Horizontal, marker.Vertical))
		}
	}

	switch mode {
	case resizeMode:
		resize()
	case fitMode:
		fit()
	case liquidMode:
		liquid()
	case cropMode:
		crop()
	case cropmidMode:
		cropmid()
	}
	watermark()

	operations = append(operations, ops.Convert{format})

	return
}

func roundedIntegerDivision(n, m int) int {
	if (n < 0) == (m < 0) {
		return (n + m/2) / m
	} else { // -5 / 6 should round to -1
		return (n - m/2) / m
	}
}

var stringToMode = map[string]mode{
	"":        resizeMode,
	"resize":  resizeMode,
	"fit":     fitMode,
	"crop":    cropMode,
	"cropmid": cropmidMode,
	"liquid":  liquidMode,
}

var supportedFormats = map[string]string{
	"":					"jpeg",
	"jpg":			"jpeg",
	"jpeg":			"jpeg",
	"gif":			"gif",
	"webp":			"webp",
	"bmp":			"bmp",
	"png":			"png",
	"tiff":			"tiff",
}

type mode int

const (
	fitMode = mode(1 + iota)
	cropMode
	cropmidMode
	liquidMode
	resizeMode

	widthParam  = "width"
	heightParam = "height"
	modeParam   = "mode"
	formatParam = "format"
	cropxParam  = "cropx"
	cropyParam  = "cropy"
	urlParam		= "url"
)

// returns validated parameters from request and error if invalid
func getParams(a *fasthttp.Args) (w, h, cropx, cropy int, mode mode, format, url string, err error) {

	if strings.Contains(a.String(), "%3F") { // %3F = ?
		err = errors.New("Invalid characters in request!")
		return
	}

	defer func() {
		if msg := recover(); msg != nil {
			err = msg.(error)
		}
	}()

	w = getUint(a, widthParam)
	h = getUint(a, heightParam)

	cropx = getUint(a, cropxParam)
	cropy = getUint(a, cropyParam)

	mode = stringToMode[string(a.Peek(modeParam))]
	if mode == 0 {
		err = errors.New("Invalid mode!")
		return
	}

	format, formatFound := supportedFormats[strings.ToLower(string(a.Peek(formatParam)))]
	if !formatFound {
		err = errors.New("Invalid format '" + string(a.Peek(formatParam)) + "'!")
		return
	}
	// TODO: verify that the format is one we support.
	// We do not want to support TXT, for instance

	url = string(a.Peek(urlParam))

	a.Del(widthParam)
	a.Del(heightParam)
	a.Del(modeParam)
	a.Del(formatParam)
	a.Del(cropxParam)
	a.Del(cropyParam)
	a.Del(urlParam)

	if a.Len() != 0 {
		err = errors.New("Invalid parameter " + a.String())
		return
	}

	err = nil
	return
}

func getUint(a *fasthttp.Args, param string) int {
	v, err := a.GetUint(param)
	if isParseError(err) {
		panic(err)
	}
	if v == -1 {
		v = 0
	}
	return v
}

func isParseError(err error) bool {
	return err != nil && err != fasthttp.ErrNoArgValue
}
<commit_msg>format<commit_after>package main

import (
	"errors"
	"github.com/phzfi/RIC/server/logging"
	"github.com/phzfi/RIC/server/ops"
	"github.com/valyala/fasthttp"
	"strings"
)

func ParseURI(uri *fasthttp.URI, source ops.ImageSource, marker ops.Watermarker) (operations []ops.Operation, format string, err, invalid error) {
	filename := string(uri.Path())

	w, h, cropx, cropy, mode, format, url, invalid := getParams(uri.QueryArgs())
	if invalid != nil {
		return
	}

	if url != "" {
		source.AddRoot(url)
	}

	ow, oh, err := source.ImageSize(filename)
	if err != nil {
		return
	}

	operations = []ops.Operation{source.LoadImageOp(filename)}

	adjustWidth := func() {
		w = roundedIntegerDivision(h*ow, oh)
	}

	adjustHeight := func() {
		h = roundedIntegerDivision(w*oh, ow)
	}

	adjustSize := func() {
		if h == 0 && w != 0 {
			adjustHeight()
		} else if h != 0 && w == 0 {
			adjustWidth()
		} else if w == 0 && h == 0 {
			w, h = ow, oh
		}
	}

	denyUpscale := func() {
		h0 := h
		w0 := w
		if w > ow {
			h = roundedIntegerDivision(ow*h0, w0)
			w = ow
		}
		if h > oh || h > h0 {
			w = roundedIntegerDivision(oh*w0, h0)
			h = oh
		}
	}

	resize := func() {
		denyUpscale()
		adjustSize()
		operations = append(operations, ops.Resize{w, h})
	}

	liquid := func() {
		denyUpscale()
		adjustSize()
		operations = append(operations, ops.LiquidRescale{w, h})
	}

	crop := func() {
		if w == 0 {
			w = ow
		}
		if h == 0 {
			h = oh
		}
		operations = append(operations, ops.Crop{w, h, cropx, cropy})
	}

	cropmid := func() {
		if w == 0 || w > ow {
			w = ow
		}
		if h == 0 || h > oh {
			h = oh
		}
		midW := roundedIntegerDivision(ow, 2)
		midH := roundedIntegerDivision(oh, 2)
		cropx := midW - roundedIntegerDivision(w, 2)
		cropy := midH - roundedIntegerDivision(h, 2)
		operations = append(operations, ops.Crop{w, h, cropx, cropy})
	}

	fit := func() {
		if w > ow {
			w = ow
		}
		if h > oh {
			h = oh
		}
		if w != 0 && h != 0 {
			if ow*h > w*oh {
				adjustHeight()
			} else {
				adjustWidth()
			}
			operations = append(operations, ops.Resize{w, h})
		} else {
			resize()
		}
	}

	watermark := func() {
		heightOK := h > marker.MinHeight && h < marker.MaxHeight
		widthOK := w > marker.MinWidth && w < marker.MaxWidth
		if marker.AddMark && heightOK && widthOK {
			logging.Debug("Adding watermarkOp")
			operations = append(operations, ops.WatermarkOp(marker.WatermarkImage, marker.Horizontal, marker.Vertical))
		}
	}

	switch mode {
	case resizeMode:
		resize()
	case fitMode:
		fit()
	case liquidMode:
		liquid()
	case cropMode:
		crop()
	case cropmidMode:
		cropmid()
	}
	watermark()

	operations = append(operations, ops.Convert{format})

	return
}

func roundedIntegerDivision(n, m int) int {
	if (n < 0) == (m < 0) {
		return (n + m/2) / m
	} else { // -5 / 6 should round to -1
		return (n - m/2) / m
	}
}

var stringToMode = map[string]mode{
	"":        resizeMode,
	"resize":  resizeMode,
	"fit":     fitMode,
	"crop":    cropMode,
	"cropmid": cropmidMode,
	"liquid":  liquidMode,
}

var supportedFormats = map[string]string{
	"":					"jpeg",
	"jpg":			"jpeg",
	"jpeg":			"jpeg",
	"gif":			"gif",
	"webp":			"webp",
	"bmp":			"bmp",
	"png":			"png",
	"tiff":			"tiff",
}

type mode int

const (
	fitMode = mode(1 + iota)
	cropMode
	cropmidMode
	liquidMode
	resizeMode

	widthParam  = "width"
	heightParam = "height"
	modeParam   = "mode"
	formatParam = "format"
	cropxParam  = "cropx"
	cropyParam  = "cropy"
	urlParam    = "url"
)

// returns validated parameters from request and error if invalid
func getParams(a *fasthttp.Args) (w, h, cropx, cropy int, mode mode, format, url string, err error) {

	if strings.Contains(a.String(), "%3F") { // %3F = ?
		err = errors.New("Invalid characters in request!")
		return
	}

	defer func() {
		if msg := recover(); msg != nil {
			err = msg.(error)
		}
	}()

	w = getUint(a, widthParam)
	h = getUint(a, heightParam)

	cropx = getUint(a, cropxParam)
	cropy = getUint(a, cropyParam)

	mode = stringToMode[string(a.Peek(modeParam))]
	if mode == 0 {
		err = errors.New("Invalid mode!")
		return
	}

	format, formatFound := supportedFormats[strings.ToLower(string(a.Peek(formatParam)))]
	if !formatFound {
		err = errors.New("Invalid format '" + string(a.Peek(formatParam)) + "'!")
		return
	}
	// TODO: verify that the format is one we support.
	// We do not want to support TXT, for instance

	url = string(a.Peek(urlParam))

	a.Del(widthParam)
	a.Del(heightParam)
	a.Del(modeParam)
	a.Del(formatParam)
	a.Del(cropxParam)
	a.Del(cropyParam)
	a.Del(urlParam)

	if a.Len() != 0 {
		err = errors.New("Invalid parameter " + a.String())
		return
	}

	err = nil
	return
}

func getUint(a *fasthttp.Args, param string) int {
	v, err := a.GetUint(param)
	if isParseError(err) {
		panic(err)
	}
	if v == -1 {
		v = 0
	}
	return v
}

func isParseError(err error) bool {
	return err != nil && err != fasthttp.ErrNoArgValue
}
<|endoftext|>


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

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
	"runtime"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//Log.Debugc(func() string {
	//	return fmt.Sprintf("(%s) Starting sender", this.cid())
	//})

	this.wgStarted.Done()
	//Log.Debugc(func() string {
	//	return fmt.Sprintf("sender_1(%s)", this.cid())
	//})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			//Log.Debugc(func() string {
			//	return fmt.Sprintf("sender_2(%s)", this.cid())
			//})
			_, err := this.out.WriteTo(conn)
			Log.Debugc(func() string {
				return fmt.Sprintf("sender_3(%s)", this.cid())
			})
			if err != nil {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_4(%s)", this.cid())
				})
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}else {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_5(%s)", this.cid())
				})
				Log.Infoc(func() string {
					return fmt.Sprintf("向conn些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
	Log.Debugc(func() string {
		return fmt.Sprintf("sender_6(%s)", this.cid())
	})
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	var b []byte
	var index int64
	var ok bool

	for i := 0; i < 99; i++ {
		b, index, ok = this.in.ReadBuffer()
		if ok {
			break
		}
		runtime.Gosched()
	}
	defer this.in.ReadCommit(index)
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (error) {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	_, err := msg.Encode(b)
	if err != nil {
		return err
	}

	for i := 0; i < 100; i++ {
		if this.out.WriteBuffer(b); {
			break
		}
		runtime.Gosched()
	}

	this.outStat.increment(int64(1))

	return nil
}
<commit_msg>修改buffer读写限制<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
	"runtime"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//Log.Debugc(func() string {
	//	return fmt.Sprintf("(%s) Starting sender", this.cid())
	//})

	this.wgStarted.Done()
	//Log.Debugc(func() string {
	//	return fmt.Sprintf("sender_1(%s)", this.cid())
	//})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			//Log.Debugc(func() string {
			//	return fmt.Sprintf("sender_2(%s)", this.cid())
			//})
			_, err := this.out.WriteTo(conn)
			Log.Debugc(func() string {
				return fmt.Sprintf("sender_3(%s)", this.cid())
			})
			if err != nil {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_4(%s)", this.cid())
				})
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}else {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_5(%s)", this.cid())
				})
				Log.Infoc(func() string {
					return fmt.Sprintf("向conn些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
	Log.Debugc(func() string {
		return fmt.Sprintf("sender_6(%s)", this.cid())
	})
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	var b []byte
	var index int64
	var ok bool

	for i := 0; i < 99; i++ {
		b, index, ok = this.in.ReadBuffer()
		if ok {
			break
		}
		runtime.Gosched()
	}
	defer this.in.ReadCommit(index)
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (error) {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	_, err := msg.Encode(b)
	if err != nil {
		return err
	}

	for i := 0; i < 100; i++ {
		if this.out.WriteBuffer(b) {
			break
		}
		runtime.Gosched()
	}

	this.outStat.increment(int64(1))

	return nil
}
<|endoftext|>


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

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
	"runtime"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//Log.Debugc(func() string {
	//	return fmt.Sprintf("(%s) Starting sender", this.cid())
	//})

	this.wgStarted.Done()
	//Log.Debugc(func() string {
	//	return fmt.Sprintf("sender_1(%s)", this.cid())
	//})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			//Log.Debugc(func() string {
			//	return fmt.Sprintf("sender_2(%s)", this.cid())
			//})
			_, err := this.out.WriteTo(conn)
			Log.Debugc(func() string {
				return fmt.Sprintf("sender_3(%s)", this.cid())
			})
			if err != nil {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_4(%s)", this.cid())
				})
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}else {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_5(%s)", this.cid())
				})
				Log.Infoc(func() string {
					return fmt.Sprintf("向conn些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
	Log.Debugc(func() string {
		return fmt.Sprintf("sender_6(%s)", this.cid())
	})
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	var b []byte
	var index int64
	var ok bool

	for i := 0; i < 99; i++ {
		if this.isDone() {
			return nil, err
		}
		b, index, ok = this.in.ReadBuffer()
		if ok {
			break
		}
		runtime.Gosched()
	}
	defer this.in.ReadCommit(index)
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (error) {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	_, err := msg.Encode(b)
	if err != nil {
		return err
	}

	for i := 0; i < 100; i++ {
		if this.isDone() {
			return 0, io.EOF
		}
		if this.out.WriteBuffer(b) {
			break
		}
		runtime.Gosched()
	}

	this.outStat.increment(int64(1))

	return nil
}
<commit_msg>修改buffer读写限制<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
	"runtime"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//Log.Debugc(func() string {
	//	return fmt.Sprintf("(%s) Starting sender", this.cid())
	//})

	this.wgStarted.Done()
	//Log.Debugc(func() string {
	//	return fmt.Sprintf("sender_1(%s)", this.cid())
	//})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			//Log.Debugc(func() string {
			//	return fmt.Sprintf("sender_2(%s)", this.cid())
			//})
			_, err := this.out.WriteTo(conn)
			Log.Debugc(func() string {
				return fmt.Sprintf("sender_3(%s)", this.cid())
			})
			if err != nil {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_4(%s)", this.cid())
				})
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}else {
				Log.Debugc(func() string {
					return fmt.Sprintf("sender_5(%s)", this.cid())
				})
				Log.Infoc(func() string {
					return fmt.Sprintf("向conn些数据成功！")
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
	Log.Debugc(func() string {
		return fmt.Sprintf("sender_6(%s)", this.cid())
	})
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	var b []byte
	var index int64
	var ok bool

	for i := 0; i < 99; i++ {
		if this.isDone() {
			return nil, err
		}
		b, index, ok = this.in.ReadBuffer()
		if ok {
			break
		}
		runtime.Gosched()
	}
	defer this.in.ReadCommit(index)
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (error) {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	_, err := msg.Encode(b)
	if err != nil {
		return err
	}

	for i := 0; i < 100; i++ {
		if this.isDone() {
			return io.EOF
		}
		if this.out.WriteBuffer(b) {
			break
		}
		runtime.Gosched()
	}

	this.outStat.increment(int64(1))

	return nil
}
<|endoftext|>


<commit_before>package service

import (
	"github.com/boivie/lovebeat/backend"
	"github.com/boivie/lovebeat/config"
	"github.com/boivie/lovebeat/metrics"
	"github.com/boivie/lovebeat/model"
	"github.com/op/go-logging"

	"regexp"
	"time"
)

type Services struct {
	be                   backend.Backend
	services             map[string]*Service
	views                map[string]View
	upsertServiceCmdChan chan *upsertServiceCmd
	deleteServiceCmdChan chan string
	getServicesChan      chan *getServicesCmd
	getServiceChan       chan *getServiceCmd
	getViewsChan         chan *getViewsCmd
	getViewChan          chan *getViewCmd
}

const (
	MAX_UNPROCESSED_PACKETS = 1000
	EXPIRY_INTERVAL         = 1
)

var (
	log      = logging.MustGetLogger("lovebeat")
	counters = metrics.NopMetrics()
	StateMap = map[string]int{
		model.STATE_PAUSED:  0,
		model.STATE_OK:      1,
		model.STATE_WARNING: 2,
		model.STATE_ERROR:   3,
	}
)

func (svcs *Services) updateView(view View, ts int64) {
	var ref = view
	view.update(ts)
	if view.data.State != ref.data.State {
		view.save(svcs.be, &ref, ts)
		if view.hasAlert(&ref) {
			// TODO: Send alert
		}
	}
	svcs.views[view.data.Name] = view
}

func (svcs *Services) updateMatchingViews(ts int64, serviceName string) {
	for _, view := range svcs.views {
		if view.contains(serviceName) {
			svcs.updateView(view, ts)
		}
	}
}

func (svcs *Services) getService(name string) *Service {
	var s, ok = svcs.services[name]
	if !ok {
		log.Debug("Asked for unknown service %s", name)
		s = newService(name)
		svcs.services[name] = s
	}
	return s
}

func (svcs *Services) Monitor(cfg config.Config) {
	period := time.Duration(EXPIRY_INTERVAL) * time.Second
	ticker := time.NewTicker(period)
	svcs.reload(cfg)

	for {
		select {
		case <-ticker.C:
			var ts = now()
			for _, s := range svcs.services {
				if s.data.State == model.STATE_PAUSED ||
					s.data.State == s.stateAt(ts) {
					continue
				}
				var ref = *s
				s.update(ts)
				s.save(svcs.be, &ref, ts)
				svcs.updateMatchingViews(ts, s.name())
			}
		case c := <-svcs.getServicesChan:
			var ret []model.Service
			var view, ok = svcs.views[c.View]
			if ok {
				for _, s := range svcs.services {
					if view.contains(s.name()) {
						ret = append(ret, s.data)
					}
				}
			}
			c.Reply <- ret
		case c := <-svcs.getServiceChan:
			var ret = svcs.services[c.Name]
			if ret == nil {
				c.Reply <- nil
			} else {
				c.Reply <- &ret.data
			}
		case c := <-svcs.getViewsChan:
			var ret []model.View
			for _, v := range svcs.views {
				ret = append(ret, v.data)
			}
			c.Reply <- ret
		case c := <-svcs.getViewChan:
			if ret, ok := svcs.views[c.Name]; ok {
				c.Reply <- &ret.data
			} else {
				c.Reply <- nil
			}
		case c := <-svcs.deleteServiceCmdChan:
			log.Info("SERVICE '%s', deleted", c)
			var ts = now()
			var s = svcs.getService(c)
			delete(svcs.services, s.name())
			svcs.be.DeleteService(s.name())
			svcs.updateMatchingViews(ts, s.name())
		case c := <-svcs.upsertServiceCmdChan:
			var ts = now()
			var s = svcs.getService(c.Service)
			var ref = *s

			if c.RegisterBeat {
				s.registerBeat(ts)
			}

			// Don't re-calculate 'auto' if we already have values
			if c.WarningTimeout == TIMEOUT_AUTO &&
				s.data.WarningTimeout == -1 {
				s.data.WarningTimeout = TIMEOUT_AUTO
				s.data.PreviousBeats = make([]int64, model.PREVIOUS_BEATS_COUNT)
			} else if c.WarningTimeout == TIMEOUT_CLEAR {
				s.data.WarningTimeout = TIMEOUT_CLEAR
			} else if c.WarningTimeout > 0 {
				s.data.WarningTimeout = c.WarningTimeout
			}
			if c.ErrorTimeout == TIMEOUT_AUTO &&
				s.data.ErrorTimeout == -1 {
				s.data.ErrorTimeout = TIMEOUT_AUTO
				s.data.PreviousBeats = make([]int64, model.PREVIOUS_BEATS_COUNT)
			} else if c.ErrorTimeout == TIMEOUT_CLEAR {
				s.data.ErrorTimeout = TIMEOUT_CLEAR
			} else if c.ErrorTimeout > 0 {
				s.data.ErrorTimeout = c.ErrorTimeout
			}
			s.update(ts)
			s.save(svcs.be, &ref, ts)
			svcs.updateMatchingViews(ts, s.name())
		}
	}
}

func (svcs *Services) createAllView() View {
	var ree, _ = regexp.Compile("")
	return View{
		services: svcs.services,
		data: model.View{
			Name:   "all",
			Regexp: "",
			State:  model.STATE_PAUSED,
		},
		ree: ree,
	}
}

func (svcs *Services) reload(cfg config.Config) {
	svcs.services = make(map[string]*Service)

	for _, s := range svcs.be.LoadServices() {
		var svc = &Service{data: *s}
		if svc.data.PreviousBeats == nil ||
			len(svc.data.PreviousBeats) != model.PREVIOUS_BEATS_COUNT {
			svc.data.PreviousBeats = make([]int64, model.PREVIOUS_BEATS_COUNT)
		}
		svcs.services[s.Name] = svc
	}

	views := make([]View, 0)
	views = append(views, svcs.createAllView())

	backendViews := svcs.be.LoadViews()

	for name, v := range cfg.Views {
		var ree, _ = regexp.Compile(v.Regexp)
		view := View{
			services: svcs.services,
			data: model.View{
				Name:   name,
				Regexp: v.Regexp,
				State:  model.STATE_PAUSED,
			},
			ree: ree,
		}
		views = append(views, view)
	}

	svcs.views = make(map[string]View)
	ts := now()
	for _, view := range views {
		if be, ok := backendViews[view.data.Name]; ok {
			view.data.State = be.State
			view.data.LastUpdated = be.LastUpdated
			view.data.IncidentNbr = be.IncidentNbr
		}
		log.Info("Created view '%s' ('%s'), state = %s",
			view.data.Name, view.data.Regexp, view.data.State)
		svcs.views[view.data.Name] = view
		svcs.updateView(view, ts)
	}
}

func NewServices(beiface backend.Backend, m metrics.Metrics) *Services {
	counters = m
	svcs := new(Services)
	svcs.be = beiface
	svcs.deleteServiceCmdChan = make(chan string, 5)
	svcs.upsertServiceCmdChan = make(chan *upsertServiceCmd, MAX_UNPROCESSED_PACKETS)
	svcs.getServicesChan = make(chan *getServicesCmd, 5)
	svcs.getServiceChan = make(chan *getServiceCmd, 5)
	svcs.getViewsChan = make(chan *getViewsCmd, 5)
	svcs.getViewChan = make(chan *getViewCmd, 5)

	return svcs
}
<commit_msg>Refactoring<commit_after>package service

import (
	"github.com/boivie/lovebeat/backend"
	"github.com/boivie/lovebeat/config"
	"github.com/boivie/lovebeat/metrics"
	"github.com/boivie/lovebeat/model"
	"github.com/op/go-logging"

	"regexp"
	"time"
)

type Services struct {
	be                   backend.Backend
	services             map[string]*Service
	views                map[string]View
	upsertServiceCmdChan chan *upsertServiceCmd
	deleteServiceCmdChan chan string
	getServicesChan      chan *getServicesCmd
	getServiceChan       chan *getServiceCmd
	getViewsChan         chan *getViewsCmd
	getViewChan          chan *getViewCmd
}

const (
	MAX_UNPROCESSED_PACKETS = 1000
	EXPIRY_INTERVAL         = 1
)

var (
	log      = logging.MustGetLogger("lovebeat")
	counters = metrics.NopMetrics()
	StateMap = map[string]int{
		model.STATE_PAUSED:  0,
		model.STATE_OK:      1,
		model.STATE_WARNING: 2,
		model.STATE_ERROR:   3,
	}
)

func (svcs *Services) updateView(view View, ts int64) {
	var ref = view
	view.update(ts)
	if view.data.State != ref.data.State {
		view.save(svcs.be, &ref, ts)
		if view.hasAlert(&ref) {
			// TODO: Send alert
		}
	}
	svcs.views[view.data.Name] = view
}

func (svcs *Services) updateMatchingViews(ts int64, serviceName string) {
	for _, view := range svcs.views {
		if view.contains(serviceName) {
			svcs.updateView(view, ts)
		}
	}
}

func (svcs *Services) getService(name string) *Service {
	var s, ok = svcs.services[name]
	if !ok {
		log.Debug("Asked for unknown service %s", name)
		s = newService(name)
		svcs.services[name] = s
	}
	return s
}

func (svcs *Services) Monitor(cfg config.Config) {
	period := time.Duration(EXPIRY_INTERVAL) * time.Second
	ticker := time.NewTicker(period)
	svcs.reload(cfg)

	for {
		select {
		case <-ticker.C:
			var ts = now()
			for _, s := range svcs.services {
				if s.data.State == model.STATE_PAUSED ||
					s.data.State == s.stateAt(ts) {
					continue
				}
				var ref = *s
				s.update(ts)
				s.save(svcs.be, &ref, ts)
				svcs.updateMatchingViews(ts, s.name())
			}
		case c := <-svcs.getServicesChan:
			var ret []model.Service
			var view, ok = svcs.views[c.View]
			if ok {
				for _, s := range svcs.services {
					if view.contains(s.name()) {
						ret = append(ret, s.data)
					}
				}
			}
			c.Reply <- ret
		case c := <-svcs.getServiceChan:
			var ret = svcs.services[c.Name]
			if ret == nil {
				c.Reply <- nil
			} else {
				c.Reply <- &ret.data
			}
		case c := <-svcs.getViewsChan:
			var ret []model.View
			for _, v := range svcs.views {
				ret = append(ret, v.data)
			}
			c.Reply <- ret
		case c := <-svcs.getViewChan:
			if ret, ok := svcs.views[c.Name]; ok {
				c.Reply <- &ret.data
			} else {
				c.Reply <- nil
			}
		case c := <-svcs.deleteServiceCmdChan:
			log.Info("SERVICE '%s', deleted", c)
			var ts = now()
			var s = svcs.getService(c)
			delete(svcs.services, s.name())
			svcs.be.DeleteService(s.name())
			svcs.updateMatchingViews(ts, s.name())
		case c := <-svcs.upsertServiceCmdChan:
			var ts = now()
			var s = svcs.getService(c.Service)
			var ref = *s

			if c.RegisterBeat {
				s.registerBeat(ts)
			}

			// Don't re-calculate 'auto' if we already have values
			if c.WarningTimeout == TIMEOUT_AUTO &&
				s.data.WarningTimeout == -1 {
				s.data.WarningTimeout = TIMEOUT_AUTO
				s.data.PreviousBeats = make([]int64, model.PREVIOUS_BEATS_COUNT)
			} else if c.WarningTimeout == TIMEOUT_CLEAR {
				s.data.WarningTimeout = TIMEOUT_CLEAR
			} else if c.WarningTimeout > 0 {
				s.data.WarningTimeout = c.WarningTimeout
			}
			if c.ErrorTimeout == TIMEOUT_AUTO &&
				s.data.ErrorTimeout == -1 {
				s.data.ErrorTimeout = TIMEOUT_AUTO
				s.data.PreviousBeats = make([]int64, model.PREVIOUS_BEATS_COUNT)
			} else if c.ErrorTimeout == TIMEOUT_CLEAR {
				s.data.ErrorTimeout = TIMEOUT_CLEAR
			} else if c.ErrorTimeout > 0 {
				s.data.ErrorTimeout = c.ErrorTimeout
			}
			s.update(ts)
			s.save(svcs.be, &ref, ts)
			svcs.updateMatchingViews(ts, s.name())
		}
	}
}

func (svcs *Services) loadViewsFromConfig(cfg config.Config) []View {
	views := make([]View, 0)

	views = append(views, View{
		services: svcs.services,
		data: model.View{
			Name:   "all",
			Regexp: "",
			State:  model.STATE_PAUSED,
		},
		ree: regexp.MustCompile(""),
	})

	for name, v := range cfg.Views {
		var ree, _ = regexp.Compile(v.Regexp)
		view := View{
			services: svcs.services,
			data: model.View{
				Name:   name,
				Regexp: v.Regexp,
				State:  model.STATE_PAUSED,
			},
			ree: ree,
		}
		views = append(views, view)
	}
	return views
}

func (svcs *Services) reload(cfg config.Config) {
	svcs.services = make(map[string]*Service)

	for _, s := range svcs.be.LoadServices() {
		var svc = &Service{data: *s}
		if svc.data.PreviousBeats == nil ||
			len(svc.data.PreviousBeats) != model.PREVIOUS_BEATS_COUNT {
			svc.data.PreviousBeats = make([]int64, model.PREVIOUS_BEATS_COUNT)
		}
		svcs.services[s.Name] = svc
	}

	views := svcs.loadViewsFromConfig(cfg)
	backendViews := svcs.be.LoadViews()

	svcs.views = make(map[string]View)
	ts := now()
	for _, view := range views {
		if be, ok := backendViews[view.data.Name]; ok {
			view.data.State = be.State
			view.data.LastUpdated = be.LastUpdated
			view.data.IncidentNbr = be.IncidentNbr
		}
		log.Info("Created view '%s' ('%s'), state = %s",
			view.data.Name, view.data.Regexp, view.data.State)
		svcs.views[view.data.Name] = view
		svcs.updateView(view, ts)
	}
}

func NewServices(beiface backend.Backend, m metrics.Metrics) *Services {
	counters = m
	svcs := new(Services)
	svcs.be = beiface
	svcs.deleteServiceCmdChan = make(chan string, 5)
	svcs.upsertServiceCmdChan = make(chan *upsertServiceCmd, MAX_UNPROCESSED_PACKETS)
	svcs.getServicesChan = make(chan *getServicesCmd, 5)
	svcs.getServiceChan = make(chan *getServiceCmd, 5)
	svcs.getViewsChan = make(chan *getViewsCmd, 5)
	svcs.getViewChan = make(chan *getViewCmd, 5)

	return svcs
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"fmt"
	"log"
	"strings"
)

type Services map[string][]string

type ServiceDirectory struct {
	services map[string]*ServiceType
	index    map[string][]string
	workers  map[string]*Worker
}

type ServiceType struct {
	name      string
	instances map[string]*Worker
}

func NewServiceMessage(serviceArray []string) (services Services, err error) {
	services = Services{}

	for _, service := range serviceArray {
		sSplit := strings.Split(service, ".")

		if len(sSplit) != 2 {
			err := errors.New("Invalid service list")
			return services, err
		}

		serviceType := sSplit[0]
		serviceInstance := sSplit[1]

		services[serviceType] = append(services[serviceType], serviceInstance)
	}

	return
}

func NewServiceDirectory() *ServiceDirectory {
	return &ServiceDirectory{
		services: make(map[string]*ServiceType),
		workers:  make(map[string]*Worker),
		index:    make(map[string][]string),
	}
}

func (serviceDirectory *ServiceDirectory) AddWorker(identity string, services Services) (err error) {
	// validate services
	name := fmt.Sprintf("%q", identity)
	invalidServices := serviceDirectory.validateServicesForWorker(name, services)

	if len(invalidServices) > 0 {
		errString := fmt.Sprintf(
			"Worker %s cannot register services %s",
			identity, strings.Join(invalidServices, ", "),
		)
		err = errors.New(errString)
		return err
	}

	// create or return worker
	worker, exists := serviceDirectory.workers[identity]

	if exists == false {
		worker = NewWorker(identity, services)

		serviceDirectory.workers[identity] = worker
	}

	for serviceTypeName, serviceInstances := range services {
		// create or return serviceType
		serviceType, exists := serviceDirectory.services[serviceTypeName]

		if exists == false {
			serviceType = &ServiceType{
				name:      serviceTypeName,
				instances: make(map[string]*Worker),
			}

			serviceDirectory.services[serviceTypeName] = serviceType
		}

		// register serviceInstance, or err if already found
		for _, serviceInstanceName := range serviceInstances {
			log.Printf("Adding worker %s to %s.%s",
				worker.Name, serviceTypeName, serviceInstanceName)
			serviceType.instances[serviceInstanceName] = worker
			// add to index
			serviceDirectory.index[serviceTypeName] = append(serviceDirectory.index[serviceTypeName], serviceInstanceName)
		}

	}
	log.Printf("?: Worker - %q", worker)

	return
}

func (serviceDirectory *ServiceDirectory) RemoveWorker(worker *Worker) {
	for sType, sInstances := range worker.Services {
		for _, sInstance := range sInstances {
			// copy the index instances array
			instances := append([]string(nil), serviceDirectory.index[sType]...)

			for i, inst := range serviceDirectory.index[sType] {
				if inst == sInstance {
					// remove matched instance from array
					instances = append(
						instances[:i],
						instances[i+1:]...,
					)
				}
			}

			if len(instances) == 0 {
				// remove entry from index
				delete(serviceDirectory.index, sType)
			} else {
				// replace index with new array
				serviceDirectory.index[sType] = instances
			}

			//remove worker reference from services
			delete(serviceDirectory.services[sType].instances, sInstance)
		}
	}

	// remove worker
	delete(serviceDirectory.workers, worker.Identity)
}

func (serviceDirectory *ServiceDirectory) validateServicesForWorker(name string, services Services) (invalidServices []string) {
	for sType, sInstances := range services {
		// if the serviceType doesnt exist, all the instances can be registered
		_, exists := serviceDirectory.index[sType]
		if exists == false {
			continue
		}

		for _, sInstance := range sInstances {
			for _, instance := range serviceDirectory.index[sType] {
				if instance == sInstance {
					invalidServices = append(
						invalidServices, fmt.Sprintf("%s.%s", sType, sInstance),
					)
				}
			}
		}
	}

	return
}

func (serviceDirectory *ServiceDirectory) WorkerForService(serviceTypeName string, serviceInstanceName string) (serviceWorker *Worker, err error) {
	// check for serviceType
	serviceType, exists := serviceDirectory.index[serviceTypeName]

	if !exists {
		err = errors.New("Unknown serviceType")
		return
	}

	log.Printf("Found serviceType %q", serviceType)

	for _, serviceInstance := range serviceType {
		if serviceInstance == serviceInstanceName {
			log.Printf("Match: %q", serviceDirectory.services[serviceTypeName])
			serviceWorker, exists = serviceDirectory.services[serviceTypeName].instances[serviceInstanceName]
			if exists {
				log.Printf("Found serviceWorker: %q", serviceWorker)
				return
			}
		}
	}

	err = errors.New("Unknown serviceInstance")
	return
}

func (serviceDirectory *ServiceDirectory) ServiceExists(serviceType string, serviceInstance string) (exists bool) {
	// check for serviceType
	services, exists := serviceDirectory.index[serviceType]

	if exists == false {
		return false
	}

	for _, si := range services {
		if si == serviceInstance {
			return true
		}
	}

	return false
}
<commit_msg>Print less rubbish to the logs<commit_after>package service

import (
	"errors"
	"fmt"
	"log"
	"strings"
)

type Services map[string][]string

type ServiceDirectory struct {
	services map[string]*ServiceType
	index    map[string][]string
	workers  map[string]*Worker
}

type ServiceType struct {
	name      string
	instances map[string]*Worker
}

func NewServiceMessage(serviceArray []string) (services Services, err error) {
	services = Services{}

	for _, service := range serviceArray {
		sSplit := strings.Split(service, ".")

		if len(sSplit) != 2 {
			err := errors.New("Invalid service list")
			return services, err
		}

		serviceType := sSplit[0]
		serviceInstance := sSplit[1]

		services[serviceType] = append(services[serviceType], serviceInstance)
	}

	return
}

func NewServiceDirectory() *ServiceDirectory {
	return &ServiceDirectory{
		services: make(map[string]*ServiceType),
		workers:  make(map[string]*Worker),
		index:    make(map[string][]string),
	}
}

func (serviceDirectory *ServiceDirectory) AddWorker(identity string, services Services) (err error) {
	// validate services
	name := fmt.Sprintf("%q", identity)
	invalidServices := serviceDirectory.validateServicesForWorker(name, services)

	if len(invalidServices) > 0 {
		errString := fmt.Sprintf(
			"Worker %s cannot register services %s",
			identity, strings.Join(invalidServices, ", "),
		)
		err = errors.New(errString)
		return err
	}

	// create or return worker
	worker, exists := serviceDirectory.workers[identity]

	if exists == false {
		worker = NewWorker(identity, services)

		serviceDirectory.workers[identity] = worker
	}

	for serviceTypeName, serviceInstances := range services {
		// create or return serviceType
		serviceType, exists := serviceDirectory.services[serviceTypeName]

		if exists == false {
			serviceType = &ServiceType{
				name:      serviceTypeName,
				instances: make(map[string]*Worker),
			}

			serviceDirectory.services[serviceTypeName] = serviceType
		}

		// register serviceInstance, or err if already found
		for _, serviceInstanceName := range serviceInstances {
			log.Printf("Adding worker %s to %s.%s",
				worker.Name, serviceTypeName, serviceInstanceName)
			serviceType.instances[serviceInstanceName] = worker
			// add to index
			serviceDirectory.index[serviceTypeName] = append(serviceDirectory.index[serviceTypeName], serviceInstanceName)
		}

	}
	log.Printf("?: Worker - %q", worker)

	return
}

func (serviceDirectory *ServiceDirectory) RemoveWorker(worker *Worker) {
	for sType, sInstances := range worker.Services {
		for _, sInstance := range sInstances {
			// copy the index instances array
			instances := append([]string(nil), serviceDirectory.index[sType]...)

			for i, inst := range serviceDirectory.index[sType] {
				if inst == sInstance {
					// remove matched instance from array
					instances = append(
						instances[:i],
						instances[i+1:]...,
					)
				}
			}

			if len(instances) == 0 {
				// remove entry from index
				delete(serviceDirectory.index, sType)
			} else {
				// replace index with new array
				serviceDirectory.index[sType] = instances
			}

			//remove worker reference from services
			delete(serviceDirectory.services[sType].instances, sInstance)
		}
	}

	// remove worker
	delete(serviceDirectory.workers, worker.Identity)
}

func (serviceDirectory *ServiceDirectory) validateServicesForWorker(name string, services Services) (invalidServices []string) {
	for sType, sInstances := range services {
		// if the serviceType doesnt exist, all the instances can be registered
		_, exists := serviceDirectory.index[sType]
		if exists == false {
			continue
		}

		for _, sInstance := range sInstances {
			for _, instance := range serviceDirectory.index[sType] {
				if instance == sInstance {
					invalidServices = append(
						invalidServices, fmt.Sprintf("%s.%s", sType, sInstance),
					)
				}
			}
		}
	}

	return
}

func (serviceDirectory *ServiceDirectory) WorkerForService(serviceTypeName string, serviceInstanceName string) (serviceWorker *Worker, err error) {
	// check for serviceType
	serviceType, exists := serviceDirectory.index[serviceTypeName]

	if !exists {
		err = errors.New("Unknown serviceType")
		return
	}

	log.Printf("Found serviceType %q", serviceType)

	for _, serviceInstance := range serviceType {
		if serviceInstance == serviceInstanceName {
			serviceWorker, exists = serviceDirectory.services[serviceTypeName].instances[serviceInstanceName]
			if exists {
				log.Printf("Found serviceWorker: %q", serviceWorker.Name)
				return
			}
		}
	}

	err = errors.New("Unknown serviceInstance")
	return
}

func (serviceDirectory *ServiceDirectory) ServiceExists(serviceType string, serviceInstance string) (exists bool) {
	// check for serviceType
	services, exists := serviceDirectory.index[serviceType]

	if exists == false {
		return false
	}

	for _, si := range services {
		if si == serviceInstance {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// Package feature package defines the names of the current feature flags.
package feature

// TODO (anastasiamac 2015-03-02)
// Features that have commands that can be blocked,
// command list for "juju block" and "juju unblock"
// needs to be maintained until we can dynamically discover
// these commands.

// LogErrorStack is a developer feature flag to have the LoggedErrorStack
// function in the utils package write out the error stack as defined by the
// errors package to the logger.  The ability to log the error stack is very
// useful in those error cases where you really don't expect there to be a
// failure.  This means that the developers with this flag set will see the
// stack trace in the log output, but normal deployments never will.
const LogErrorStack = "log-error-stack"

// LegacyUpstart is used to indicate that the version-based init system
// discovery code (service.VersionInitSystem) should return upstart
// instead of systemd for vivid and newer.
const LegacyUpstart = "legacy-upstart"

// PostNetCLIMVP is used to indicate that additional networking
// commands will be available in addition to the networking MVP ones
// (space list|create, subnet list|add).
const PostNetCLIMVP = "post-net-cli-mvp"

// ImageMetadata allows custom image metadata to be recorded in state.
const ImageMetadata = "image-metadata"

// DeveloperMode allows access to developer specific commands and behaviour.
const DeveloperMode = "developer-mode"

// CrossModelRelations allows cross model relations functionality.
const CrossModelRelations = "cross-model"

// LXDStorage enables the LXD storage provider.
const LXDStorage = "lxd-storage"

// PersistentStorage enables the persistent storage feature.
const PersistentStorage = "persistent-storage"

// StrictMigration will cause migration to error if there are unexported
// values for annotations, status, status history, or settings.
const StrictMigration = "strict-migration"
<commit_msg>Add feature flag for CAAS models<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// Package feature package defines the names of the current feature flags.
package feature

// TODO (anastasiamac 2015-03-02)
// Features that have commands that can be blocked,
// command list for "juju block" and "juju unblock"
// needs to be maintained until we can dynamically discover
// these commands.

// LogErrorStack is a developer feature flag to have the LoggedErrorStack
// function in the utils package write out the error stack as defined by the
// errors package to the logger.  The ability to log the error stack is very
// useful in those error cases where you really don't expect there to be a
// failure.  This means that the developers with this flag set will see the
// stack trace in the log output, but normal deployments never will.
const LogErrorStack = "log-error-stack"

// LegacyUpstart is used to indicate that the version-based init system
// discovery code (service.VersionInitSystem) should return upstart
// instead of systemd for vivid and newer.
const LegacyUpstart = "legacy-upstart"

// PostNetCLIMVP is used to indicate that additional networking
// commands will be available in addition to the networking MVP ones
// (space list|create, subnet list|add).
const PostNetCLIMVP = "post-net-cli-mvp"

// ImageMetadata allows custom image metadata to be recorded in state.
const ImageMetadata = "image-metadata"

// DeveloperMode allows access to developer specific commands and behaviour.
const DeveloperMode = "developer-mode"

// CrossModelRelations allows cross model relations functionality.
const CrossModelRelations = "cross-model"

// LXDStorage enables the LXD storage provider.
const LXDStorage = "lxd-storage"

// PersistentStorage enables the persistent storage feature.
const PersistentStorage = "persistent-storage"

// StrictMigration will cause migration to error if there are unexported
// values for annotations, status, status history, or settings.
const StrictMigration = "strict-migration"

// CAASModel enables creating models on CAAS infrastructure (k8s, etc)
const CAASModel = "caas-model"
<|endoftext|>


<commit_before>package log

import (
	"github.com/getconversio/go-utils/util"
	"github.com/puddingfactory/logentrus"
	log "github.com/sirupsen/logrus"
)

// Setup options for logging with logrus
type Options struct {
	Environment     string
	LogentriesToken string
}

func Setup(options *Options) {
	if options == nil {
		options = &Options{}
	}

	level := log.InfoLevel

	if options.Environment == "development" {
		level = log.DebugLevel
		log.SetFormatter(&log.TextFormatter{ForceColors: true})
	}

	log.SetLevel(level)

	if options.LogentriesToken != "" {
		hook, err := logentrus.New(options.LogentriesToken, &logentrus.Opts{Priority: level})
		util.PanicOnError("Could not create logentries hook", err)
		log.AddHook(hook)
	}
}
<commit_msg>Update logentrus repo url (#12)<commit_after>package log

import (
	"github.com/getconversio/go-utils/util"
	"github.com/jonathan-robertson/logentrus"
	log "github.com/sirupsen/logrus"
)

// Setup options for logging with logrus
type Options struct {
	Environment     string
	LogentriesToken string
}

func Setup(options *Options) {
	if options == nil {
		options = &Options{}
	}

	level := log.InfoLevel

	if options.Environment == "development" {
		level = log.DebugLevel
		log.SetFormatter(&log.TextFormatter{ForceColors: true})
	}

	log.SetLevel(level)

	if options.LogentriesToken != "" {
		hook, err := logentrus.New(options.LogentriesToken, &logentrus.Opts{Priority: level})
		util.PanicOnError("Could not create logentries hook", err)
		log.AddHook(hook)
	}
}
<|endoftext|>


<commit_before>package sqls

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"

	"fmt"
	"os"
	"testing"
)

var (
	mysqlConn *sql.DB
	username  string = "root"
	passwd    string = "123456"
)

func init() {
	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(127.0.0.1:3306)/golib_test", username, passwd))
	if err != nil {
		fmt.Printf("connect to mysql error\n")
		os.Exit(1)
	}

	mysqlConn = db
}

func TestCreateTable(t *testing.T) {
	db := &SqlWrap{db: mysqlConn}
	sql := "CREATE TABLE IF NOT EXISTS `users` (`id` bigint unsigned NOT NULL AUTO_INCREMENT, `name` varchar(255) NOT NULL, PRIMARY KEY (`id`) ) ENGINE=InnoDB DEFAULT CHARSET=utf8"
	_, err := db.Exec(sql)
	if err != nil {
		t.Fatalf("create table error, %v", err)
	}
}

func TestInsertOne(t *testing.T) {
	db := &SqlWrap{db: mysqlConn}
	sql := "INSERT INTO users (name) VALUES ('test')"
	_, err := db.Exec(sql)
	if err != nil {
		t.Fatalf("insert into db error, %v", err)
	}
}

func TestInsertTx(t *testing.T) {
	db := &SqlWrap{db: mysqlConn}
	sql := "INSERT INTO users (name) VALUES (?)"

	tx, err := db.Begin()
	if err != nil {
		t.Fatalf("start tx error: %v", err)
	}

	for i := 0; i < 5; i++ {
		_, err := db.ExecTx(tx, sql, fmt.Sprintf("test%d", i))
		if err != nil {
			db.Rollback(tx)
			t.Fatalf("tx insert into db error, %v", err)
		}
	}

	db.Commit(tx)
}
<commit_msg>sql: fix mysql passwd for travis-ci<commit_after>package sqls

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"

	"fmt"
	"os"
	"testing"
)

var (
	mysqlConn *sql.DB
	username  string = "root"
	passwd    string = ""
)

func init() {
	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(127.0.0.1:3306)/golib_test", username, passwd))
	if err != nil {
		fmt.Printf("connect to mysql error\n")
		os.Exit(1)
	}

	mysqlConn = db
}

func TestCreateTable(t *testing.T) {
	db := &SqlWrap{db: mysqlConn}
	sql := "CREATE TABLE IF NOT EXISTS `users` (`id` bigint unsigned NOT NULL AUTO_INCREMENT, `name` varchar(255) NOT NULL, PRIMARY KEY (`id`) ) ENGINE=InnoDB DEFAULT CHARSET=utf8"
	_, err := db.Exec(sql)
	if err != nil {
		t.Fatalf("create table error, %v", err)
	}
}

func TestInsertOne(t *testing.T) {
	db := &SqlWrap{db: mysqlConn}
	sql := "INSERT INTO users (name) VALUES ('test')"
	_, err := db.Exec(sql)
	if err != nil {
		t.Fatalf("insert into db error, %v", err)
	}
}

func TestInsertTx(t *testing.T) {
	db := &SqlWrap{db: mysqlConn}
	sql := "INSERT INTO users (name) VALUES (?)"

	tx, err := db.Begin()
	if err != nil {
		t.Fatalf("start tx error: %v", err)
	}

	for i := 0; i < 5; i++ {
		_, err := db.ExecTx(tx, sql, fmt.Sprintf("test%d", i))
		if err != nil {
			db.Rollback(tx)
			t.Fatalf("tx insert into db error, %v", err)
		}
	}

	db.Commit(tx)
}
<|endoftext|>


<commit_before>package migrations

import "github.com/BurntSushi/migration"

func RemoveWorkerIds(tx migration.LimitedTx) error {
	_, err := tx.Exec(`
		ALTER TABLE volumes ADD COLUMN worker_name text;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		UPDATE volumes v set worker_name =
		(SELECT name from workers w where w.id = v.worker_id);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		BEGIN WORK;
		LOCK TABLE volumes IN ACCESS EXCLUSIVE MODE;
		DELETE FROM volumes WHERE worker_name IS NULL;
		ALTER TABLE volumes ALTER COLUMN worker_name SET NOT NULL;
		COMMIT WORK;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		CREATE INDEX volumes_worker_name ON volumes (worker_name);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
    ALTER TABLE volumes DROP CONSTRAINT volumes_worker_id_handle_key;
    ALTER TABLE volumes ADD UNIQUE (worker_name, handle);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE volumes DROP COLUMN worker_id;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE containers
		ADD COLUMN worker_name TEXT NULL;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		UPDATE containers c SET worker_name =
		(select name from workers w where w.id = c.worker_id);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE containers
		DROP COLUMN worker_id;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE workers
		DROP COLUMN id;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE workers
		ADD PRIMARY KEY (name);
	`)
	return err
}
<commit_msg>can't lock on table during parallel test run<commit_after>package migrations

import "github.com/BurntSushi/migration"

func RemoveWorkerIds(tx migration.LimitedTx) error {
	_, err := tx.Exec(`
		ALTER TABLE volumes ADD COLUMN worker_name text;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		UPDATE volumes v set worker_name =
		(SELECT name from workers w where w.id = v.worker_id);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		DELETE FROM volumes WHERE worker_name IS NULL;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE volumes ALTER COLUMN worker_name SET NOT NULL;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		CREATE INDEX volumes_worker_name ON volumes (worker_name);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
    ALTER TABLE volumes DROP CONSTRAINT volumes_worker_id_handle_key;
    ALTER TABLE volumes ADD UNIQUE (worker_name, handle);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE volumes DROP COLUMN worker_id;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE containers
		ADD COLUMN worker_name TEXT NULL;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		UPDATE containers c SET worker_name =
		(select name from workers w where w.id = c.worker_id);
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE containers
		DROP COLUMN worker_id;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE workers
		DROP COLUMN id;
	`)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		ALTER TABLE workers
		ADD PRIMARY KEY (name);
	`)
	return err
}
<|endoftext|>


<commit_before>package app_files

import (
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/gofileutils/fileutils"
)

type AppFiles interface {
	AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error)
	CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) (err error)
	CountFiles(directory string) int64
	WalkAppFiles(dir string, onEachFile func(string, string) error) (err error)
}

type ApplicationFiles struct{}

func (appfiles ApplicationFiles) AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error) {
	dir, err = filepath.Abs(dir)
	if err != nil {
		return
	}

	err = appfiles.WalkAppFiles(dir, func(fileName string, fullPath string) error {
		fileInfo, err := os.Lstat(fullPath)
		if err != nil {
			return err
		}

		appFile := models.AppFileFields{
			Path: filepath.ToSlash(fileName),
			Size: fileInfo.Size(),
		}

		if fileInfo.IsDir() {
			appFile.Sha1 = "0"
			appFile.Size = 0
		} else {
			hash := sha1.New()
			file, err := os.Open(fullPath)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(hash, file)
			if err != nil {
				return err
			}

			appFile.Sha1 = fmt.Sprintf("%x", hash.Sum(nil))
		}

		appFiles = append(appFiles, appFile)

		return nil
	})

	return
}

func (appfiles ApplicationFiles) CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) error {
	for _, file := range appFiles {
		fromPath := filepath.Join(fromDir, file.Path)
		srcFileInfo, err := os.Stat(fromPath)
		if err != nil {
			return err
		}

		toPath := filepath.Join(toDir, file.Path)

		if srcFileInfo.IsDir() {
			err = os.MkdirAll(toPath, srcFileInfo.Mode())
			if err != nil {
				return err
			}
			continue
		}

		var dst *os.File
		dst, err = fileutils.Create(toPath)
		if err != nil {
			return err
		}
		defer dst.Close()

		dst.Chmod(srcFileInfo.Mode())

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

		_, err = io.Copy(dst, src)
		if err != nil {
			return err
		}
	}

	return nil
}

func (appfiles ApplicationFiles) CountFiles(directory string) int64 {
	var count int64
	appfiles.WalkAppFiles(directory, func(_, _ string) error {
		count++
		return nil
	})
	return count
}

func (appfiles ApplicationFiles) WalkAppFiles(dir string, onEachFile func(string, string) error) (err error) {
	cfIgnore := loadIgnoreFile(dir)
	walkFunc := func(fullPath string, f os.FileInfo, inErr error) (err error) {
		err = inErr
		if err != nil {
			return
		}

		if fullPath == dir {
			return
		}

		if !f.Mode().IsRegular() && !f.IsDir() {
			return
		}

		fileRelativePath, _ := filepath.Rel(dir, fullPath)
		fileRelativeUnixPath := filepath.ToSlash(fileRelativePath)

		if !cfIgnore.FileShouldBeIgnored(fileRelativeUnixPath) {
			err = onEachFile(fileRelativePath, fullPath)
		}

		return
	}

	err = filepath.Walk(dir, walkFunc)
	return
}

func loadIgnoreFile(dir string) CfIgnore {
	fileContents, err := ioutil.ReadFile(filepath.Join(dir, ".cfignore"))
	if err != nil {
		return NewCfIgnore("")
	}

	return NewCfIgnore(string(fileContents))
}
<commit_msg>Defer close of files in CopyFiles properly<commit_after>package app_files

import (
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/gofileutils/fileutils"
)

type AppFiles interface {
	AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error)
	CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) (err error)
	CountFiles(directory string) int64
	WalkAppFiles(dir string, onEachFile func(string, string) error) (err error)
}

type ApplicationFiles struct{}

func (appfiles ApplicationFiles) AppFilesInDir(dir string) (appFiles []models.AppFileFields, err error) {
	dir, err = filepath.Abs(dir)
	if err != nil {
		return
	}

	err = appfiles.WalkAppFiles(dir, func(fileName string, fullPath string) error {
		fileInfo, err := os.Lstat(fullPath)
		if err != nil {
			return err
		}

		appFile := models.AppFileFields{
			Path: filepath.ToSlash(fileName),
			Size: fileInfo.Size(),
		}

		if fileInfo.IsDir() {
			appFile.Sha1 = "0"
			appFile.Size = 0
		} else {
			hash := sha1.New()
			file, err := os.Open(fullPath)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(hash, file)
			if err != nil {
				return err
			}

			appFile.Sha1 = fmt.Sprintf("%x", hash.Sum(nil))
		}

		appFiles = append(appFiles, appFile)

		return nil
	})

	return
}

func (appfiles ApplicationFiles) CopyFiles(appFiles []models.AppFileFields, fromDir, toDir string) error {
	for _, file := range appFiles {
		err := func() error {
			fromPath := filepath.Join(fromDir, file.Path)
			srcFileInfo, err := os.Stat(fromPath)
			if err != nil {
				return err
			}

			toPath := filepath.Join(toDir, file.Path)

			if srcFileInfo.IsDir() {
				err = os.MkdirAll(toPath, srcFileInfo.Mode())
				if err != nil {
					return err
				}
				return nil
			}

			var dst *os.File
			dst, err = fileutils.Create(toPath)
			if err != nil {
				return err
			}
			defer dst.Close()

			dst.Chmod(srcFileInfo.Mode())

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

			_, err = io.Copy(dst, src)
			if err != nil {
				return err
			}

			return nil
		}()

		if err != nil {
			return err
		}
	}

	return nil
}

func (appfiles ApplicationFiles) CountFiles(directory string) int64 {
	var count int64
	appfiles.WalkAppFiles(directory, func(_, _ string) error {
		count++
		return nil
	})
	return count
}

func (appfiles ApplicationFiles) WalkAppFiles(dir string, onEachFile func(string, string) error) (err error) {
	cfIgnore := loadIgnoreFile(dir)
	walkFunc := func(fullPath string, f os.FileInfo, inErr error) (err error) {
		err = inErr
		if err != nil {
			return
		}

		if fullPath == dir {
			return
		}

		if !f.Mode().IsRegular() && !f.IsDir() {
			return
		}

		fileRelativePath, _ := filepath.Rel(dir, fullPath)
		fileRelativeUnixPath := filepath.ToSlash(fileRelativePath)

		if !cfIgnore.FileShouldBeIgnored(fileRelativeUnixPath) {
			err = onEachFile(fileRelativePath, fullPath)
		}

		return
	}

	err = filepath.Walk(dir, walkFunc)
	return
}

func loadIgnoreFile(dir string) CfIgnore {
	fileContents, err := ioutil.ReadFile(filepath.Join(dir, ".cfignore"))
	if err != nil {
		return NewCfIgnore("")
	}

	return NewCfIgnore(string(fileContents))
}
<|endoftext|>


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

package auth

import (
	"io/ioutil"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fsnotify/fsnotify"
)

var (
	// TokenFileName is the file in which we store the current token
	TokenFileName = "auth.token"

	currentToken    string
	currentIdentity Identity
	zerotime        time.Time
	expiration      time.Time
	cond            = &sync.Cond{L: &sync.Mutex{}}
)

func now() time.Time {
	return time.Now().UTC()
}

func expired() bool {
	if currentToken == "" {
		return true
	}
	if expiration.IsZero() {
		return false
	}
	return expiration.Add(-expirationDelta).Before(now())
}

const (
	// expirationDelta is a margin of error during which a token should be
	// considered expired. This should help avoid expiration races when server
	// times don't match
	expirationDelta = 10 * time.Second
)

// TokenFunc is a function that can return an authentication token and its
// expiration time
type TokenFunc func() (string, int64, error)

// RefreshToken gets a new token, sets it as the current, and returns the expiration time
func RefreshToken(f TokenFunc, filename string) (int64, error) {
	log.Debug("Refreshing authentication token")
	token, expires, err := f()
	if err != nil {
		return 0, err
	}
	updateToken(token, time.Unix(expires, 0), filename)
	log.WithField("expiration", expires).Info("Received new authentication token")
	return expires, err
}

// AuthToken returns an unexpired auth token, blocking if necessary until
// authenticated
func AuthToken() string {
	cond.L.Lock()
	defer cond.L.Unlock()
	for expired() {
		cond.Wait()
	}
	return currentToken
}

func TokenLoop(f TokenFunc, tokenfile string, done chan interface{}) {
	for {
		expires, err := RefreshToken(f, tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to obtain authentication token. Retrying in 10s")
			select {
			case <-done:
				return
			case <-time.After(10 * time.Second):
			}
			continue
		}
		// Reauthenticate 1 minute before the token expires
		expiration := time.Unix(expires, 0).Sub(now())
		refresh := expiration - time.Duration(1*time.Minute)
		select {
		case <-done:
			return
		case <-time.After(refresh):
		}
	}
}

// Watch a token file for changes. Load the token when those changes occur.
func WatchTokenFile(tokenfile string, done <-chan interface{}) error {
	log := log.WithFields(logrus.Fields{
		"tokenfile": tokenfile,
	})

	loadToken := func() {
		data, err := ioutil.ReadFile(tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to load authentication token from file. Continuing to watch for changes")
		}
		// No need to handle expires or save file, because we're loading from the file rather
		// than re-requesting authentication tokens
		updateToken(string(data), zerotime, "")
		log.Infof("Updated authentication token from disk")
	}

	// An initial token load without any file changes
	loadToken()

	// Now watch for changes
	filechanges, err := NotifyOnChange(tokenfile, fsnotify.Write|fsnotify.Create, done)
	if err != nil {
		return err
	}
	for _ = range filechanges {
		loadToken()
	}
	return nil
}

func updateToken(token string, expires time.Time, filename string) {
	cond.L.Lock()
	currentToken = token
	currentIdentity = getIdentityFromToken(token)
	expiration = expires
	if filename != "" {
		ioutil.WriteFile(filename, []byte(token), 0600)
	}
	cond.L.Unlock()
	cond.Broadcast()
}

func getIdentityFromToken(token string) Identity {
	identity, err := ParseJWTIdentity(token)
	if err != nil {
		log.WithError(err).Error("Unable to obtain identity from token.")
		return nil
	}
	return identity
}

//Return the current identity or nil if the token is not yet available
func CurrentIdentity() Identity {
	cond.L.Lock()
	defer cond.L.Unlock()
	return currentIdentity
}
<commit_msg>Addressed comments<commit_after>// Copyright 2016 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package auth

import (
	"io/ioutil"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fsnotify/fsnotify"
)

const (
	// expirationDelta is a margin of error during which a token should be
	// considered expired. This should help avoid expiration races when server
	// times don't match
	expirationDelta = 10 * time.Second
)

var (
	// TokenFileName is the file in which we store the current token
	TokenFileName = "auth.token"

	currentToken    string
	currentIdentity Identity
	zerotime        time.Time
	expiration      time.Time
	cond            = &sync.Cond{L: &sync.Mutex{}}
)

// TokenFunc is a function that can return an authentication token and its
// expiration time
type TokenFunc func() (string, int64, error)

// RefreshToken gets a new token, sets it as the current, and returns the expiration time
func RefreshToken(f TokenFunc, filename string) (int64, error) {
	log.Debug("Refreshing authentication token")
	token, expires, err := f()
	if err != nil {
		return 0, err
	}
	updateToken(token, time.Unix(expires, 0), filename)
	log.WithField("expiration", expires).Info("Received new authentication token")
	return expires, err
}

// AuthToken returns an unexpired auth token, blocking if necessary until
// authenticated
func AuthToken() string {
	cond.L.Lock()
	defer cond.L.Unlock()
	for expired() {
		cond.Wait()
	}
	return currentToken
}

// CurrentIdentity returns the identity represented by the currently-live token,
// or nil if the token is not yet available
func CurrentIdentity() Identity {
	cond.L.Lock()
	defer cond.L.Unlock()
	return currentIdentity
}

// TokenLoop accepts a function that returns an expiring token. It will then
// periodically refresh that token, one minute before it is due to expire,
// setting the result as the current live token, until the done channel is
// closed.
func TokenLoop(f TokenFunc, tokenfile string, done chan interface{}) {
	for {
		expires, err := RefreshToken(f, tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to obtain authentication token. Retrying in 10s")
			select {
			case <-done:
				return
			case <-time.After(10 * time.Second):
			}
			continue
		}
		// Reauthenticate 1 minute before the token expires
		expiration := time.Unix(expires, 0).Sub(now())
		refresh := expiration - time.Duration(1*time.Minute)
		select {
		case <-done:
			return
		case <-time.After(refresh):
		}
	}
}

// Watch a token file for changes. Load the token when those changes occur.
func WatchTokenFile(tokenfile string, done <-chan interface{}) error {
	log := log.WithFields(logrus.Fields{
		"tokenfile": tokenfile,
	})

	loadToken := func() {
		data, err := ioutil.ReadFile(tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to load authentication token from file. Continuing to watch for changes")
		}
		// No need to handle expires or save file, because we're loading from the file rather
		// than re-requesting authentication tokens
		updateToken(string(data), zerotime, "")
		log.Infof("Updated authentication token from disk")
	}

	// An initial token load without any file changes
	loadToken()

	// Now watch for changes
	filechanges, err := NotifyOnChange(tokenfile, fsnotify.Write|fsnotify.Create, done)
	if err != nil {
		return err
	}
	for _ = range filechanges {
		loadToken()
	}
	return nil
}

func now() time.Time {
	return time.Now().UTC()
}

// expired returns whether the currently-live token has expired. If the token
// is empty, is is considered expired for these purposes. If it is the zero
// instant, it never expires. Otherwise, it is expired if the expiration time
// minus a margin of error is in the past.
func expired() bool {
	if currentToken == "" {
		return true
	}
	if expiration.IsZero() {
		return false
	}
	return expiration.Add(-expirationDelta).Before(now())
}

func updateToken(token string, expires time.Time, filename string) {
	cond.L.Lock()
	currentToken = token
	currentIdentity = getIdentityFromToken(token)
	expiration = expires
	if filename != "" {
		ioutil.WriteFile(filename, []byte(token), 0600)
	}
	cond.L.Unlock()
	cond.Broadcast()
}

func getIdentityFromToken(token string) Identity {
	identity, err := ParseJWTIdentity(token)
	if err != nil {
		log.WithError(err).Error("Unable to obtain identity from token.")
		return nil
	}
	return identity
}
<|endoftext|>