<commit_before>package server

import (
	"errors"

	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/commands"
	"github.com/emersion/go-imap/responses"
)

// imap errors in Authenticated state.
var (
	ErrNotAuthenticated = errors.New("Not authenticated")
)

type Select struct {
	commands.Select
}

func (cmd *Select) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	items := []string{
		imap.MailboxFlags, imap.MailboxPermanentFlags,
		imap.MailboxMessages, imap.MailboxRecent, imap.MailboxUnseen,
		imap.MailboxUidNext, imap.MailboxUidValidity,
	}

	status, err := mbox.Status(items)
	if err != nil {
		return err
	}

	ctx.Mailbox = mbox
	ctx.MailboxReadOnly = cmd.ReadOnly || status.ReadOnly

	res := &responses.Select{Mailbox: status}
	if err := conn.WriteResp(res); err != nil {
		return err
	}

	code := imap.CodeReadWrite
	if ctx.MailboxReadOnly {
		code = imap.CodeReadOnly
	}
	return ErrStatusResp(&imap.StatusResp{
		Type: imap.StatusOk,
		Code: code,
	})
}

type Create struct {
	commands.Create
}

func (cmd *Create) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	return ctx.User.CreateMailbox(cmd.Mailbox)
}

type Delete struct {
	commands.Delete
}

func (cmd *Delete) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	return ctx.User.DeleteMailbox(cmd.Mailbox)
}

type Rename struct {
	commands.Rename
}

func (cmd *Rename) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	return ctx.User.RenameMailbox(cmd.Existing, cmd.New)
}

type Subscribe struct {
	commands.Subscribe
}

func (cmd *Subscribe) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	return mbox.Subscribe()
}

type Unsubscribe struct {
	commands.Unsubscribe
}

func (cmd *Unsubscribe) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	return mbox.Unsubscribe()
}

type List struct {
	commands.List
}

func (cmd *List) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	ch := make(chan *imap.MailboxInfo)
	res := &responses.List{Mailboxes: ch, Subscribed: cmd.Subscribed}

	done := make(chan error, 1)
	go (func() {
		done <- conn.WriteResp(res)
		close(done)
	})()

	mailboxes, err := ctx.User.ListMailboxes(cmd.Subscribed)
	if err != nil {
		close(ch)
		return err
	}

	for _, mbox := range mailboxes {
		info, err := mbox.Info()
		if err != nil {
			close(ch)
			return err
		}

		// An empty ("" string) mailbox name argument is a special request to return
		// the hierarchy delimiter and the root name of the name given in the
		// reference.
		if cmd.Mailbox == "" {
			ch <- &imap.MailboxInfo{
				Attributes: []string{imap.NoSelectAttr},
				Delimiter:  info.Delimiter,
				Name:       info.Delimiter,
			}
			break
		}

		if info.Match(cmd.Reference, cmd.Mailbox) {
			ch <- info
		}
	}

	close(ch)

	return <-done
}

type Status struct {
	commands.Status
}

func (cmd *Status) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	status, err := mbox.Status(cmd.Items)
	if err != nil {
		return err
	}

	res := &responses.Status{Mailbox: status}
	return conn.WriteResp(res)
}

type Append struct {
	commands.Append
}

func (cmd *Append) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		// TODO: add [TRYCREATE] to the NO response
		return err
	}

	if err := mbox.CreateMessage(cmd.Flags, cmd.Date, cmd.Message); err != nil {
		return err
	}

	// If APPEND targets the currently selected mailbox, send an untagged EXISTS
	// Do this only if the backend doesn't send updates itself
	if conn.Server().Updates == nil && ctx.Mailbox != nil && ctx.Mailbox.Name() == mbox.Name() {
		status, err := mbox.Status([]string{imap.MailboxMessages})
		if err != nil {
			return err
		}

		res := &responses.Select{Mailbox: status}
		if err := conn.WriteResp(res); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>server: filter out unrequested fields when sending STATUS response<commit_after>package server

import (
	"errors"

	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/commands"
	"github.com/emersion/go-imap/responses"
)

// imap errors in Authenticated state.
var (
	ErrNotAuthenticated = errors.New("Not authenticated")
)

type Select struct {
	commands.Select
}

func (cmd *Select) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	items := []string{
		imap.MailboxFlags, imap.MailboxPermanentFlags,
		imap.MailboxMessages, imap.MailboxRecent, imap.MailboxUnseen,
		imap.MailboxUidNext, imap.MailboxUidValidity,
	}

	status, err := mbox.Status(items)
	if err != nil {
		return err
	}

	ctx.Mailbox = mbox
	ctx.MailboxReadOnly = cmd.ReadOnly || status.ReadOnly

	res := &responses.Select{Mailbox: status}
	if err := conn.WriteResp(res); err != nil {
		return err
	}

	code := imap.CodeReadWrite
	if ctx.MailboxReadOnly {
		code = imap.CodeReadOnly
	}
	return ErrStatusResp(&imap.StatusResp{
		Type: imap.StatusOk,
		Code: code,
	})
}

type Create struct {
	commands.Create
}

func (cmd *Create) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	return ctx.User.CreateMailbox(cmd.Mailbox)
}

type Delete struct {
	commands.Delete
}

func (cmd *Delete) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	return ctx.User.DeleteMailbox(cmd.Mailbox)
}

type Rename struct {
	commands.Rename
}

func (cmd *Rename) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	return ctx.User.RenameMailbox(cmd.Existing, cmd.New)
}

type Subscribe struct {
	commands.Subscribe
}

func (cmd *Subscribe) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	return mbox.Subscribe()
}

type Unsubscribe struct {
	commands.Unsubscribe
}

func (cmd *Unsubscribe) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	return mbox.Unsubscribe()
}

type List struct {
	commands.List
}

func (cmd *List) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	ch := make(chan *imap.MailboxInfo)
	res := &responses.List{Mailboxes: ch, Subscribed: cmd.Subscribed}

	done := make(chan error, 1)
	go (func() {
		done <- conn.WriteResp(res)
		close(done)
	})()

	mailboxes, err := ctx.User.ListMailboxes(cmd.Subscribed)
	if err != nil {
		close(ch)
		return err
	}

	for _, mbox := range mailboxes {
		info, err := mbox.Info()
		if err != nil {
			close(ch)
			return err
		}

		// An empty ("" string) mailbox name argument is a special request to return
		// the hierarchy delimiter and the root name of the name given in the
		// reference.
		if cmd.Mailbox == "" {
			ch <- &imap.MailboxInfo{
				Attributes: []string{imap.NoSelectAttr},
				Delimiter:  info.Delimiter,
				Name:       info.Delimiter,
			}
			break
		}

		if info.Match(cmd.Reference, cmd.Mailbox) {
			ch <- info
		}
	}

	close(ch)

	return <-done
}

type Status struct {
	commands.Status
}

func (cmd *Status) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		return err
	}

	status, err := mbox.Status(cmd.Items)
	if err != nil {
		return err
	}

	// Only keep items thqat have been requested
	items := make(map[string]interface{})
	for _, k := range cmd.Items {
		items[k] = status.Items[k]
	}
	status.Items = items

	res := &responses.Status{Mailbox: status}
	return conn.WriteResp(res)
}

type Append struct {
	commands.Append
}

func (cmd *Append) Handle(conn Conn) error {
	ctx := conn.Context()
	if ctx.User == nil {
		return ErrNotAuthenticated
	}

	mbox, err := ctx.User.GetMailbox(cmd.Mailbox)
	if err != nil {
		// TODO: add [TRYCREATE] to the NO response
		return err
	}

	if err := mbox.CreateMessage(cmd.Flags, cmd.Date, cmd.Message); err != nil {
		return err
	}

	// If APPEND targets the currently selected mailbox, send an untagged EXISTS
	// Do this only if the backend doesn't send updates itself
	if conn.Server().Updates == nil && ctx.Mailbox != nil && ctx.Mailbox.Name() == mbox.Name() {
		status, err := mbox.Status([]string{imap.MailboxMessages})
		if err != nil {
			return err
		}

		res := &responses.Select{Mailbox: status}
		if err := conn.WriteResp(res); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// 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 (
	"reflect"
	"sync"
	"time"

	"fmt"

	"github.com/jennal/goplay/aop"
	"github.com/jennal/goplay/encode"
	"github.com/jennal/goplay/filter"
	"github.com/jennal/goplay/filter/heartbeat"
	"github.com/jennal/goplay/helpers"
	"github.com/jennal/goplay/log"
	"github.com/jennal/goplay/pkg"
	"github.com/jennal/goplay/router"
	"github.com/jennal/goplay/session"
	"github.com/jennal/goplay/transfer"
)

const (
	REQUEST_TIMEOUT = 3 * time.Second
)

type requestCallbacks struct {
	startTime      time.Time
	successCallbak *Method
	failCallback   *Method
}

type ServiceClient struct {
	*session.Session

	router  *router.Router
	filters []filter.IFilter

	requestCbsMutex sync.Mutex
	requestCbs      map[pkg.PackageIDType]*requestCallbacks

	pushCbsMutex sync.Mutex
	pushCbs      map[string][]*Method
}

func NewServiceClient(cli transfer.IClient) *ServiceClient {
	result := &ServiceClient{
		Session: session.NewSession(cli),

		router: nil,
		filters: []filter.IFilter{
			heartbeat.NewHeartBeatManager(),
		},

		requestCbs: make(map[pkg.PackageIDType]*requestCallbacks),
		pushCbs:    make(map[string][]*Method),
	}
	result.setupEventLoop()

	return result
}

func (self *ServiceClient) SetRouter(router *router.Router) {
	self.router = router
}

func (self *ServiceClient) RegistFilter(filter filter.IFilter) {
	self.filters = append(self.filters, filter)
}

func (self *ServiceClient) SetFilters(filters []filter.IFilter) {
	self.filters = filters
}

func (s *ServiceClient) checkTimeoutLoop() {
	for {
		if !s.IsConnected() {
			break
		}

		ids := []pkg.PackageIDType{}

		s.requestCbsMutex.Lock()
		for id, item := range s.requestCbs {
			if time.Since(item.startTime) > REQUEST_TIMEOUT {
				ids = append(ids, id)
				item.failCallback.Call(pkg.NewErrorMessage(pkg.STAT_ERR_TIMEOUT, "Request Timeout"))
			}
		}

		for _, id := range ids {
			delete(s.requestCbs, id)
		}
		s.requestCbsMutex.Unlock()

		time.Sleep(REQUEST_TIMEOUT)
	}
}

func (s *ServiceClient) setupEventLoop() {
	s.AddListener(ON_SERVICE_DOWN, func(ok bool) {
		// log.Log(ON_SERVICE_DOWN)
		s.Disconnect()
	})

	exitChan := make(chan int, 1)
	s.On(transfer.EVENT_CLIENT_CONNECTED, s, func(client transfer.IClient) {
		if s.filters != nil && len(s.filters) > 0 {
			for _, filter := range s.filters {
				if !filter.OnNewClient(s.Session) {
					return
				}
			}
		}

		go s.checkTimeoutLoop()
		go func() {
			aop.Recover(func() {
			Loop:
				for {
					select {
					case <-exitChan:
						break Loop
					default:
						sess := session.NewSession(s)
						sess.Bind(s.ID)
						header, bodyBuf, err := sess.Recv()
						if err != nil {
							log.Errorf("Recv:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
							sess.Disconnect()
							break Loop
						}

						if header.Type != pkg.PKG_HEARTBEAT && header.Type != pkg.PKG_HEARTBEAT_RESPONSE {
							log.Logf("Recv:\n\theader => %#v\n\tbody => %#v | %v\n\terr => %v\n", header, bodyBuf, string(bodyBuf), err)
						}

						if (header.Type & pkg.PKG_RPC) == pkg.PKG_RPC {
							sess.BindClientID(header.ClientID)
							s.BindClientID(header.ClientID) /* FIXME: this may not be good enough */
						}

						//filters
						if s.filters != nil && len(s.filters) > 0 {
							for _, filter := range s.filters {
								if !filter.OnRecv(s.Session, header, bodyBuf) {
									goto Loop
								}
							}
						}

						switch header.Type {
						case pkg.PKG_REQUEST, pkg.PKG_RPC_REQUEST:
							if s.router != nil {
								results, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
								// fmt.Printf(" => Loop result: %#v\n", results)
								err = s.response(sess, header, results)
								if err != nil {
									log.Errorf("Response:\n\terr => %v\n\theader => %#v\n\tresults => %#v", err, header, results)
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_NOTIFY, pkg.PKG_RPC_NOTIFY:
							if s.router != nil {
								_, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_PUSH, pkg.PKG_RPC_PUSH:
							s.recvPush(header, bodyBuf)
						case pkg.PKG_RESPONSE, pkg.PKG_RPC_RESPONSE:
							s.recvResponse(header, bodyBuf)
						case pkg.PKG_HEARTBEAT, pkg.PKG_HEARTBEAT_RESPONSE:
							fallthrough
						default:
							log.Errorf("Can't reach here!!\n\terr => %v\n\theader => %#v\n\tbody => %#v", err, header, bodyBuf)
							break
						}
					}
				}
			}, func(err interface{}) {
				if e, ok := err.(error); ok {
					log.Error(e)
				} else if err != nil {
					log.Errorf("%#v", err)
				} else {
					log.Errorf("%#v", err)
				}

				s.Disconnect()
			})
		}()
	})
	s.Once(transfer.EVENT_CLIENT_DISCONNECTED, s, func(cli transfer.IClient) {
		exitChan <- 1
	})
}

func (s *ServiceClient) callRouteFunc(sess *session.Session, header *pkg.Header, bodyBuf []byte) ([]interface{}, error) {
	/*
	 * 1. find route func
	 * 2. unmarshal data
	 * 3. call route func
	 */
	method := s.router.Get(header.Route)
	if method == nil {
		return nil, log.NewErrorf("Can't find method with route: %s", header.Route)
	}
	val := method.NewArg(2)
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))
	decoder := encode.GetEncodeDecoder(header.Encoding)
	err := decoder.Unmarshal(bodyBuf, val)
	if err != nil {
		return nil, log.NewErrorf("Service.callRouteFunc decoder.Unmarshal failed: %v", err)
	}
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))

	var result []interface{}
	aop.Recover(func() {
		result = method.Call(sess, helpers.GetValueFromPtr(val))
	}, func(e interface{}) {
		err = e.(error)
	})

	return result, err
}

func (s *ServiceClient) response(sess *session.Session, header *pkg.Header, results []interface{}) error {
	respHeader := *header
	if header.Type == pkg.PKG_RPC_REQUEST {
		respHeader.Type = pkg.PKG_RPC_RESPONSE
	} else {
		respHeader.Type = pkg.PKG_RESPONSE
	}

	if results == nil || len(results) <= 0 {
		return sess.Send(&respHeader, []byte{})
	}

	result := results[0]
	/* check error != nil */
	if len(results) == 2 && !reflect.ValueOf(results[1]).IsNil() {
		respHeader.Status = pkg.STAT_ERR
		result = results[1]
	}

	// fmt.Println("result:", result)

	encoder := encode.GetEncodeDecoder(header.Encoding)
	body, err := encoder.Marshal(result)
	if err != nil {
		return err
	}

	return sess.Send(&respHeader, body)
}

func (s *ServiceClient) recvPush(header *pkg.Header, body []byte) {
	s.pushCbsMutex.Lock()
	list, ok := s.pushCbs[header.Route]
	s.pushCbsMutex.Unlock()

	if !ok {
		return
	}

	for _, item := range list {
		val := item.NewArg(0)
		s.Encoder.Unmarshal(body, val)
		// log.Log("==========>\t", i, "\t", val)
		item.Call(helpers.GetValueFromPtr(val))
	}
}

func (s *ServiceClient) recvResponse(header *pkg.Header, body []byte) {
	s.requestCbsMutex.Lock()
	cbs, ok := s.requestCbs[header.ID]
	if ok {
		delete(s.requestCbs, header.ID)
	}
	s.requestCbsMutex.Unlock()

	if !ok {
		return
	}

	// log.Logf("%v %v %v", header.Status, body, string(body))
	if header.Status == pkg.STAT_OK {
		val := cbs.successCallbak.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.successCallbak.Call(helpers.GetValueFromPtr(val))
			return
		}
	} else {
		val := cbs.failCallback.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.failCallback.Call(helpers.GetValueFromPtr(val))
			return
		}
	}

	cbs.failCallback.Call(pkg.NewErrorMessage(
		pkg.STAT_ERR_DECODE_FAILED,
		fmt.Sprintf("decode body failed: %#v | %v", body, string(body))))
}

func (s *ServiceClient) Request(route string, data interface{}, succCb interface{}, failCb func(*pkg.ErrorMessage)) error {
	header := s.NewHeader(pkg.PKG_RPC_REQUEST, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	cbs := requestCallbacks{
		successCallbak: NewMethod(succCb),
		failCallback:   NewMethod(failCb),
		startTime:      time.Now(),
	}

	s.requestCbsMutex.Lock()
	s.requestCbs[header.ID] = &cbs
	s.requestCbsMutex.Unlock()

	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Notify(route string, data interface{}) error {
	header := s.NewHeader(pkg.PKG_RPC_NOTIFY, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Push(route string, data interface{}) error {
	return s.Session.Push(route, data)
}

func (s *ServiceClient) AddListener(route string, callback interface{}) {
	s.pushCbsMutex.Lock()
	defer s.pushCbsMutex.Unlock()

	list, ok := s.pushCbs[route]
	if !ok {
		list = make([]*Method, 0)
		s.pushCbs[route] = list
	}

	s.pushCbs[route] = append(s.pushCbs[route], NewMethod(callback))
}
<commit_msg>change recv & disconnect to ServiceClient<commit_after>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// 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 (
	"reflect"
	"sync"
	"time"

	"fmt"

	"github.com/jennal/goplay/aop"
	"github.com/jennal/goplay/encode"
	"github.com/jennal/goplay/filter"
	"github.com/jennal/goplay/filter/heartbeat"
	"github.com/jennal/goplay/helpers"
	"github.com/jennal/goplay/log"
	"github.com/jennal/goplay/pkg"
	"github.com/jennal/goplay/router"
	"github.com/jennal/goplay/session"
	"github.com/jennal/goplay/transfer"
)

const (
	REQUEST_TIMEOUT = 3 * time.Second
)

type requestCallbacks struct {
	startTime      time.Time
	successCallbak *Method
	failCallback   *Method
}

type ServiceClient struct {
	*session.Session

	router  *router.Router
	filters []filter.IFilter

	requestCbsMutex sync.Mutex
	requestCbs      map[pkg.PackageIDType]*requestCallbacks

	pushCbsMutex sync.Mutex
	pushCbs      map[string][]*Method
}

func NewServiceClient(cli transfer.IClient) *ServiceClient {
	result := &ServiceClient{
		Session: session.NewSession(cli),

		router: nil,
		filters: []filter.IFilter{
			heartbeat.NewHeartBeatManager(),
		},

		requestCbs: make(map[pkg.PackageIDType]*requestCallbacks),
		pushCbs:    make(map[string][]*Method),
	}
	result.setupEventLoop()

	return result
}

func (self *ServiceClient) SetRouter(router *router.Router) {
	self.router = router
}

func (self *ServiceClient) RegistFilter(filter filter.IFilter) {
	self.filters = append(self.filters, filter)
}

func (self *ServiceClient) SetFilters(filters []filter.IFilter) {
	self.filters = filters
}

func (s *ServiceClient) checkTimeoutLoop() {
	for {
		if !s.IsConnected() {
			break
		}

		ids := []pkg.PackageIDType{}

		s.requestCbsMutex.Lock()
		for id, item := range s.requestCbs {
			if time.Since(item.startTime) > REQUEST_TIMEOUT {
				ids = append(ids, id)
				item.failCallback.Call(pkg.NewErrorMessage(pkg.STAT_ERR_TIMEOUT, "Request Timeout"))
			}
		}

		for _, id := range ids {
			delete(s.requestCbs, id)
		}
		s.requestCbsMutex.Unlock()

		time.Sleep(REQUEST_TIMEOUT)
	}
}

func (s *ServiceClient) setupEventLoop() {
	s.AddListener(ON_SERVICE_DOWN, func(ok bool) {
		// log.Log(ON_SERVICE_DOWN)
		s.Disconnect()
	})

	exitChan := make(chan int, 1)
	s.On(transfer.EVENT_CLIENT_CONNECTED, s, func(client transfer.IClient) {
		if s.filters != nil && len(s.filters) > 0 {
			for _, filter := range s.filters {
				if !filter.OnNewClient(s.Session) {
					return
				}
			}
		}

		go s.checkTimeoutLoop()
		go func() {
			aop.Recover(func() {
			Loop:
				for {
					select {
					case <-exitChan:
						break Loop
					default:
						sess := session.NewSession(s)
						sess.Bind(s.ID)
						header, bodyBuf, err := s.Recv()
						if err != nil {
							log.Errorf("Recv:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
							s.Disconnect()
							break Loop
						}

						// if header.Type != pkg.PKG_HEARTBEAT && header.Type != pkg.PKG_HEARTBEAT_RESPONSE {
						// 	log.Logf("Recv:\n\theader => %#v\n\tbody => %#v | %v\n\terr => %v\n", header, bodyBuf, string(bodyBuf), err)
						// }

						if (header.Type & pkg.PKG_RPC) == pkg.PKG_RPC {
							sess.BindClientID(header.ClientID)
							s.BindClientID(header.ClientID) /* FIXME: this may not be good enough */
						}

						//filters
						if s.filters != nil && len(s.filters) > 0 {
							for _, filter := range s.filters {
								if !filter.OnRecv(s.Session, header, bodyBuf) {
									goto Loop
								}
							}
						}

						switch header.Type {
						case pkg.PKG_REQUEST, pkg.PKG_RPC_REQUEST:
							if s.router != nil {
								results, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
								// fmt.Printf(" => Loop result: %#v\n", results)
								err = s.response(sess, header, results)
								if err != nil {
									log.Errorf("Response:\n\terr => %v\n\theader => %#v\n\tresults => %#v", err, header, results)
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_NOTIFY, pkg.PKG_RPC_NOTIFY:
							if s.router != nil {
								_, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_PUSH, pkg.PKG_RPC_PUSH:
							s.recvPush(header, bodyBuf)
						case pkg.PKG_RESPONSE, pkg.PKG_RPC_RESPONSE:
							s.recvResponse(header, bodyBuf)
						case pkg.PKG_HEARTBEAT, pkg.PKG_HEARTBEAT_RESPONSE:
							fallthrough
						default:
							log.Errorf("Can't reach here!!\n\terr => %v\n\theader => %#v\n\tbody => %#v", err, header, bodyBuf)
							break
						}
					}
				}
			}, func(err interface{}) {
				if e, ok := err.(error); ok {
					log.Error(e)
				} else if err != nil {
					log.Errorf("%#v", err)
				} else {
					log.Errorf("%#v", err)
				}

				s.Disconnect()
			})
		}()
	})
	s.Once(transfer.EVENT_CLIENT_DISCONNECTED, s, func(cli transfer.IClient) {
		exitChan <- 1
	})
}

func (s *ServiceClient) callRouteFunc(sess *session.Session, header *pkg.Header, bodyBuf []byte) ([]interface{}, error) {
	/*
	 * 1. find route func
	 * 2. unmarshal data
	 * 3. call route func
	 */
	method := s.router.Get(header.Route)
	if method == nil {
		return nil, log.NewErrorf("Can't find method with route: %s", header.Route)
	}
	val := method.NewArg(2)
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))
	decoder := encode.GetEncodeDecoder(header.Encoding)
	err := decoder.Unmarshal(bodyBuf, val)
	if err != nil {
		return nil, log.NewErrorf("Service.callRouteFunc decoder.Unmarshal failed: %v", err)
	}
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))

	var result []interface{}
	aop.Recover(func() {
		result = method.Call(sess, helpers.GetValueFromPtr(val))
	}, func(e interface{}) {
		err = e.(error)
	})

	return result, err
}

func (s *ServiceClient) response(sess *session.Session, header *pkg.Header, results []interface{}) error {
	respHeader := *header
	if header.Type == pkg.PKG_RPC_REQUEST {
		respHeader.Type = pkg.PKG_RPC_RESPONSE
	} else {
		respHeader.Type = pkg.PKG_RESPONSE
	}

	if results == nil || len(results) <= 0 {
		return sess.Send(&respHeader, []byte{})
	}

	result := results[0]
	/* check error != nil */
	if len(results) == 2 && !reflect.ValueOf(results[1]).IsNil() {
		respHeader.Status = pkg.STAT_ERR
		result = results[1]
	}

	// fmt.Println("result:", result)

	encoder := encode.GetEncodeDecoder(header.Encoding)
	body, err := encoder.Marshal(result)
	if err != nil {
		return err
	}

	return sess.Send(&respHeader, body)
}

func (s *ServiceClient) recvPush(header *pkg.Header, body []byte) {
	s.pushCbsMutex.Lock()
	list, ok := s.pushCbs[header.Route]
	s.pushCbsMutex.Unlock()

	if !ok {
		return
	}

	for _, item := range list {
		val := item.NewArg(0)
		s.Encoder.Unmarshal(body, val)
		// log.Log("==========>\t", i, "\t", val)
		item.Call(helpers.GetValueFromPtr(val))
	}
}

func (s *ServiceClient) recvResponse(header *pkg.Header, body []byte) {
	s.requestCbsMutex.Lock()
	cbs, ok := s.requestCbs[header.ID]
	if ok {
		delete(s.requestCbs, header.ID)
	}
	s.requestCbsMutex.Unlock()

	if !ok {
		return
	}

	// log.Logf("%v %v %v", header.Status, body, string(body))
	if header.Status == pkg.STAT_OK {
		val := cbs.successCallbak.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.successCallbak.Call(helpers.GetValueFromPtr(val))
			return
		}
	} else {
		val := cbs.failCallback.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.failCallback.Call(helpers.GetValueFromPtr(val))
			return
		}
	}

	cbs.failCallback.Call(pkg.NewErrorMessage(
		pkg.STAT_ERR_DECODE_FAILED,
		fmt.Sprintf("decode body failed: %#v | %v", body, string(body))))
}

func (s *ServiceClient) Request(route string, data interface{}, succCb interface{}, failCb func(*pkg.ErrorMessage)) error {
	header := s.NewHeader(pkg.PKG_RPC_REQUEST, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	cbs := requestCallbacks{
		successCallbak: NewMethod(succCb),
		failCallback:   NewMethod(failCb),
		startTime:      time.Now(),
	}

	s.requestCbsMutex.Lock()
	s.requestCbs[header.ID] = &cbs
	s.requestCbsMutex.Unlock()

	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Notify(route string, data interface{}) error {
	header := s.NewHeader(pkg.PKG_RPC_NOTIFY, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Push(route string, data interface{}) error {
	return s.Session.Push(route, data)
}

func (s *ServiceClient) AddListener(route string, callback interface{}) {
	s.pushCbsMutex.Lock()
	defer s.pushCbsMutex.Unlock()

	list, ok := s.pushCbs[route]
	if !ok {
		list = make([]*Method, 0)
		s.pushCbs[route] = list
	}

	s.pushCbs[route] = append(s.pushCbs[route], NewMethod(callback))
}
<|endoftext|>


<commit_before>package bigly

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/biogo/hts/sam"
)

// SkipBase holds the byte used for a Skip Operation.
var SkipBase byte = '.'

func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}

func max(a, b int) int {
	if a < b {
		return b
	}
	return a
}

// Options holds information about what is reported in the Pile
type Options struct {
	MinBaseQuality    uint8  `arg:"-q,help:base quality threshold"`
	MinMappingQuality uint8  `arg:"-Q,help:mapping quality threshold"`
	ExcludeFlag       uint16 `arg:"-F"`
	IncludeFlag       uint16 `arg:"-f"`
	MinClipLength     int    `arg:"-c,help:only count H/S clips of at least this length"`
	IncludeBases      bool   `arg:"-b,help:output each base and base quality score"`
	SplitterVerbosity int    `arg:"-s,help:0-only count; 1:count and single most frequent; 2:all SAs; 3:dont shorten positions"`
	ConcordantCutoff  int    `arg:"-o,help:distance beyond which mates are called discordant"`
}

// Pile holds the information about a single base.
type Pile struct {
	Chrom                 string
	Pos                   int
	Depth                 int    // count of reads passing filters.
	RefBase               byte   // Reference base a this position.
	MisMatches            uint32 // number of mismatches .
	ProperPairs           int    // count of reads with paired flag
	SoftStarts            uint32 // counts of base preceding an 'S' cigar op
	SoftEnds              uint32 // ...  following ...
	HardStarts            uint32 // counts of base preceding an 'H' cigar op
	HardEnds              uint32
	InsertionStarts       uint32 // counts of base preceding an 'I' cigar op
	InsertionEnds         uint32
	Deletions             uint32  // counts of deletions 'D' at this base
	Heads                 uint32  // counts of starts of reads at this base
	Tails                 uint32  // counts of ends of reads at this base
	Splitters             uint32  // count of non-secondary reads with SA tags.
	Splitters1            uint32  // count of non-secondary reads with exactly 1 SA tag.
	Bases                 []byte  // All bases from reads covering this position
	Quals                 []uint8 // All quals from reads covering this position
	InsertSizeLPs         []int32 // All insert size for left-most of pair
	InsertSizeRMs         []int32 // ...                 right-most of pair
	OrientationPlusPlus   uint32  // Paired reads mapped in +/+ orientation
	OrientationMinusMinus uint32  // Paired reads mapped in -/- orientation
	OrientationMinusPlus  uint32  // Paired reads mapped in -/+ orientation
	// count of reads where splitters were in different orientation than read.
	// only set if SplitterVerbosity > 1
	OrientationSplitter uint32
	/*
		TODO: figure out how to do this because these are excluded by ExcludeFlag.
			   could try to calculate and it will ony be non-zero if these are not excluded.C
		Duplicates             uint32  // reads counted as duplicates
		Supplementary          uint32  // reads counted as supplementary
	*/
	Discordant             uint32  // Number of reads with insert size > ConcordantCutoff
	DiscordantChrom        uint32  // Number of reads mapping on different chroms
	DiscordantChromEntropy float32 // high value means all discordants came from same chrom.
	GC65                   uint32
	GC257                  uint32
	Duplicity65            float32 // measure of lack of sequence entropy.
	Duplicity257           float32 // measure of lack of sequence entropy.
	SplitterPositions      []Position
}

// from biogo/hts
func formatQual(q []byte) []byte {
	for _, v := range q {
		if v != 0xff {
			a := make([]byte, len(q))
			for i, p := range q {
				a[i] = p + 33
			}
			return a
		}
	}
	return []byte{}
}

func formatInts(a []int) string {
	s := ""
	for _, v := range a {
		s += strconv.Itoa(v) + ","
	}
	if len(s) > 0 {
		return s[:len(s)-2]
	}
	return s
}

func mean(arr []int32) int {
	if len(arr) == 0 {
		return 0
	}
	var s float64
	for _, a := range arr {
		s += float64(a)
	}
	return int(0.5 + s/float64(len(arr)))
}

// TabString prints a tab-delimited version of the Pile
func (p Pile) TabString(o Options) string {
	spl := ""
	if o.SplitterVerbosity > 0 && len(p.SplitterPositions) > 0 {
		if o.SplitterVerbosity == 1 {
			m, c := pMode(p.SplitterPositions)
			spl = fmt.Sprintf("%d/%d/%d", m, c, len(p.SplitterPositions))
		} else {
			aspl := make([]string, len(p.SplitterPositions))
			for i, s := range p.SplitterPositions {
				aspl[i] = s.String()
			}
			spl = strings.Join(aspl, ",")
		}
	}
	return fmt.Sprintf("%s\t%d\t%d\t%c\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%.2f\t%d\t%d\t%.2f\t%.2f\t%s",
		p.Chrom,
		p.Pos+1, p.Depth, p.RefBase, p.MisMatches,
		p.ProperPairs, p.SoftStarts, p.SoftEnds,
		p.HardStarts, p.HardEnds, p.InsertionStarts, p.InsertionEnds, p.Deletions,
		//p.Heads, p.Tails,
		p.Splitters,
		p.Splitters1,
		// string(p.RefBase), string(p.Bases), string(formatQual(p.Quals)),
		mean(p.InsertSizeLPs),
		mean(p.InsertSizeRMs),
		p.OrientationPlusPlus+p.OrientationMinusPlus+p.OrientationMinusMinus+p.OrientationSplitter,
		p.Discordant,
		p.DiscordantChrom,
		p.DiscordantChromEntropy,
		p.GC65,
		p.GC257,
		p.Duplicity65,
		p.Duplicity257,
		spl,
	)
}

func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

// Update the Pile with info from the Alignment if it meets the requirements in Options
func (p *Pile) Update(o Options, alns []*Align) {
	// the Total values hold the sum of 1 / val so we can calc harmonic Mean
	// with less susceptiblity to outliers.
	var discMates []int
	for _, a := range alns {
		if a.MapQ < o.MinMappingQuality {
			continue
		}
		s := a.At(p.Pos)
		if s == nil || s.Qual < o.MinBaseQuality {
			continue
		}

		if a.Flags&sam.Paired == sam.Paired {
			if a.Flags&sam.ProperPair == sam.ProperPair {
				p.ProperPairs++
			}
			/*
				if a.Flags&sam.ProperPair != sam.ProperPair {
					p.Discordant++
				}
			*/
			// This gives much cleaner results for eat least 1 case than using the actual
			// flag
			if abs(a.Start()-a.MatePos) > o.ConcordantCutoff {
				p.Discordant++
			}
			if a.MateRef.ID() != a.Ref.ID() {
				p.DiscordantChrom++
				discMates = append(discMates, a.MateRef.ID())
			} else {
				// same chromosome.
				if a.Start() < a.MatePos && a.Flags&sam.Reverse != sam.Reverse {
					p.InsertSizeLPs = append(p.InsertSizeLPs, int32(a.MatePos-a.Start()))
				} else if a.Start() > a.MatePos && a.Flags&sam.Reverse == sam.Reverse {
					p.InsertSizeRMs = append(p.InsertSizeRMs, int32(float64(a.Start()-a.MatePos)))
				} else {
					if a.Flags&sam.Reverse == sam.Reverse && a.Flags&sam.MateReverse == sam.MateReverse {
						p.OrientationMinusMinus++
					} else if 0 == a.Flags&sam.Reverse && 0 == a.Flags&sam.MateReverse {
						p.OrientationPlusPlus++
					} else {
						p.OrientationMinusPlus++
					}
				}
			}

		}
		if a.Flags&sam.Secondary == 0 {
			if tags, ok := a.Record.Tag([]byte{'S', 'A'}); ok {
				p.Splitters++
				if bytes.Count([]byte(tags), []byte{';'}) <= 1 {
					p.Splitters1++
				}

				p.updateSplitters(o, tags, a.Strand() != -1)
			}
		}

		p.Depth++

		switch s.Right.Type() {
		case sam.CigarMatch:
			break
		case sam.CigarInsertion:
			p.InsertionStarts++
		case sam.CigarSoftClipped:
			if s.Right.Len() >= o.MinClipLength {
				p.SoftStarts++
			}
		case sam.CigarHardClipped:
			if s.Right.Len() >= o.MinClipLength {
				p.HardStarts++
			}
		}

		switch s.Left.Type() {
		case sam.CigarMatch:
			break
		case sam.CigarInsertion:
			p.InsertionEnds++
		case sam.CigarSoftClipped:
			if s.Left.Len() >= o.MinClipLength {
				p.SoftEnds++
			}
		case sam.CigarHardClipped:
			if s.Left.Len() >= o.MinClipLength {
				p.HardEnds++
			}
		}

		if s.Head {
			p.Heads++
		} else if s.Tail {
			p.Tails++
		}

		if s.At.Type() == sam.CigarDeletion {
			p.Deletions++
		}
		if o.IncludeBases {
			p.Bases = append(p.Bases, s.Base)
			p.Quals = append(p.Quals, s.Qual)
		}
		if s.Base != p.RefBase {
			p.MisMatches++
		}
	}

	if p.DiscordantChrom > 1 {
		p.DiscordantChromEntropy = float32(entropy(discMates))
	}
	// don't set this if we don't know the reference base.
	if p.RefBase == 'N' {
		p.MisMatches = 0
	}
}

// track the actual positions of the splitters and check the orientation.
func (p *Pile) updateSplitters(o Options, tags []byte, readStrand bool) {
	if o.SplitterVerbosity == 0 {
		return
	}
	sas := ParseSAs(tags)
	for _, sa := range sas {
		if sa.MapQ >= o.MinMappingQuality {
			p.SplitterPositions = append(p.SplitterPositions, Position{Chrom: string(sa.Chrom), Start: sa.Pos, End: sa.End(), Strand: sa.Strand})
			if readStrand != sa.Strand {
				p.OrientationSplitter++
			}
		}
	}
}

// Align is a sam.Record with a cursor to track position in the read and reference.
type Align struct {
	*sam.Record
	// track the into the Cigar Ops.
	CursorCigar int
	// track basepairs into the Reference
	CursorPos int
	// track basepairs into the read
	CursorRead int

	// Sequence holds the expanded sequence.
	Sequence []byte
	// track that we're alwasy moving forward.
	lastPos int
}

// CigarSummary is a summary of the context of a position for 1 alignment
type CigarSummary struct {
	At        sam.CigarOp
	Left      sam.CigarOp
	Right     sam.CigarOp
	Qual      uint8
	Head      bool
	Tail      bool
	Base      byte
	Insertion []byte
}

// At returns the CigarOp for a particular genomic position of the given read.
func (a *Align) At(pos0 int) *CigarSummary {
	if pos0+1 <= a.lastPos {
		log.Fatal("can't use align on the same position or lower position's that previous calls")
	}
	a.lastPos = pos0 + 1
	pos := a.Pos + a.CursorPos
	if pos0 < pos || len(a.Cigar) == 0 {
		return nil
	}

	if a.Sequence == nil {
		a.Sequence = a.Seq.Expand()
	}

	res := &CigarSummary{Left: a.Cigar[max(a.CursorCigar-1, 0)]}
	res.Head = pos0 == pos && a.CursorPos == 0

	for _, co := range a.Cigar[a.CursorCigar:] {
		t := co.Type()
		con := t.Consumes()
		l := co.Len()
		lr := l * con.Reference
		lq := l * con.Query

		// we can update the Left if the current position is still less
		// then the requested position.
		if pos0 > pos {
			res.Left = co
		}
		if pos+lr > pos0 {
			res.At = co
			right := co
			res.Right = co
			// add distance into read. Subtract how far before la we are.
			readi := a.CursorRead + (pos0 - pos)
			if lq > 0 {
				// if we consume bases from the read, then grab them.
				res.Base = a.Sequence[readi]
				res.Qual = a.Qual[readi]
			} else {
				res.Base = '*'
				if co.Type() == sam.CigarSkipped {
					res.Base = SkipBase
				}
			}
			if pos+lr-1 == pos0 {
				idx := a.CursorCigar + 1
				if idx == len(a.Cigar) {
					res.Tail = true
				} else {
					res.Right = a.Cigar[idx]
					right = a.Cigar[idx]
				}
			}
			if res.Right.Type() == sam.CigarInsertion {
				res.Insertion = a.Sequence[readi+1 : readi+1+right.Len()]
			}

			return res
		}
		res.Left = co
		a.CursorCigar++
		a.CursorPos += lr
		a.CursorRead += lq
		pos += lr
	}
	return nil
}
<commit_msg>only count splitter orientation changes once<commit_after>package bigly

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/biogo/hts/sam"
)

// SkipBase holds the byte used for a Skip Operation.
var SkipBase byte = '.'

func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}

func max(a, b int) int {
	if a < b {
		return b
	}
	return a
}

// Options holds information about what is reported in the Pile
type Options struct {
	MinBaseQuality    uint8  `arg:"-q,help:base quality threshold"`
	MinMappingQuality uint8  `arg:"-Q,help:mapping quality threshold"`
	ExcludeFlag       uint16 `arg:"-F"`
	IncludeFlag       uint16 `arg:"-f"`
	MinClipLength     int    `arg:"-c,help:only count H/S clips of at least this length"`
	IncludeBases      bool   `arg:"-b,help:output each base and base quality score"`
	SplitterVerbosity int    `arg:"-s,help:0-only count; 1:count and single most frequent; 2:all SAs; 3:dont shorten positions"`
	ConcordantCutoff  int    `arg:"-o,help:distance beyond which mates are called discordant"`
}

// Pile holds the information about a single base.
type Pile struct {
	Chrom                 string
	Pos                   int
	Depth                 int    // count of reads passing filters.
	RefBase               byte   // Reference base a this position.
	MisMatches            uint32 // number of mismatches .
	ProperPairs           int    // count of reads with paired flag
	SoftStarts            uint32 // counts of base preceding an 'S' cigar op
	SoftEnds              uint32 // ...  following ...
	HardStarts            uint32 // counts of base preceding an 'H' cigar op
	HardEnds              uint32
	InsertionStarts       uint32 // counts of base preceding an 'I' cigar op
	InsertionEnds         uint32
	Deletions             uint32  // counts of deletions 'D' at this base
	Heads                 uint32  // counts of starts of reads at this base
	Tails                 uint32  // counts of ends of reads at this base
	Splitters             uint32  // count of non-secondary reads with SA tags.
	Splitters1            uint32  // count of non-secondary reads with exactly 1 SA tag.
	Bases                 []byte  // All bases from reads covering this position
	Quals                 []uint8 // All quals from reads covering this position
	InsertSizeLPs         []int32 // All insert size for left-most of pair
	InsertSizeRMs         []int32 // ...                 right-most of pair
	OrientationPlusPlus   uint32  // Paired reads mapped in +/+ orientation
	OrientationMinusMinus uint32  // Paired reads mapped in -/- orientation
	OrientationMinusPlus  uint32  // Paired reads mapped in -/+ orientation
	// count of reads where splitters were in different orientation than read.
	// only set if SplitterVerbosity > 1
	OrientationSplitter uint32
	/*
		TODO: figure out how to do this because these are excluded by ExcludeFlag.
			   could try to calculate and it will ony be non-zero if these are not excluded.C
		Duplicates             uint32  // reads counted as duplicates
		Supplementary          uint32  // reads counted as supplementary
	*/
	Discordant             uint32  // Number of reads with insert size > ConcordantCutoff
	DiscordantChrom        uint32  // Number of reads mapping on different chroms
	DiscordantChromEntropy float32 // high value means all discordants came from same chrom.
	GC65                   uint32
	GC257                  uint32
	Duplicity65            float32 // measure of lack of sequence entropy.
	Duplicity257           float32 // measure of lack of sequence entropy.
	SplitterPositions      []Position
}

// from biogo/hts
func formatQual(q []byte) []byte {
	for _, v := range q {
		if v != 0xff {
			a := make([]byte, len(q))
			for i, p := range q {
				a[i] = p + 33
			}
			return a
		}
	}
	return []byte{}
}

func formatInts(a []int) string {
	s := ""
	for _, v := range a {
		s += strconv.Itoa(v) + ","
	}
	if len(s) > 0 {
		return s[:len(s)-2]
	}
	return s
}

func mean(arr []int32) int {
	if len(arr) == 0 {
		return 0
	}
	var s float64
	for _, a := range arr {
		s += float64(a)
	}
	return int(0.5 + s/float64(len(arr)))
}

// TabString prints a tab-delimited version of the Pile
func (p Pile) TabString(o Options) string {
	spl := ""
	if o.SplitterVerbosity > 0 && len(p.SplitterPositions) > 0 {
		if o.SplitterVerbosity == 1 {
			m, c := pMode(p.SplitterPositions)
			spl = fmt.Sprintf("%d/%d/%d", m, c, len(p.SplitterPositions))
		} else {
			aspl := make([]string, len(p.SplitterPositions))
			for i, s := range p.SplitterPositions {
				aspl[i] = s.String()
			}
			spl = strings.Join(aspl, ",")
		}
	}
	return fmt.Sprintf("%s\t%d\t%d\t%c\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%d\t%.2f\t%d\t%d\t%.2f\t%.2f\t%s",
		p.Chrom,
		p.Pos+1, p.Depth, p.RefBase, p.MisMatches,
		p.ProperPairs, p.SoftStarts, p.SoftEnds,
		p.HardStarts, p.HardEnds, p.InsertionStarts, p.InsertionEnds, p.Deletions,
		//p.Heads, p.Tails,
		p.Splitters,
		p.Splitters1,
		// string(p.RefBase), string(p.Bases), string(formatQual(p.Quals)),
		mean(p.InsertSizeLPs),
		mean(p.InsertSizeRMs),
		p.OrientationPlusPlus+p.OrientationMinusPlus+p.OrientationMinusMinus+p.OrientationSplitter,
		p.Discordant,
		p.DiscordantChrom,
		p.DiscordantChromEntropy,
		p.GC65,
		p.GC257,
		p.Duplicity65,
		p.Duplicity257,
		spl,
	)
}

func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

// Update the Pile with info from the Alignment if it meets the requirements in Options
func (p *Pile) Update(o Options, alns []*Align) {
	// the Total values hold the sum of 1 / val so we can calc harmonic Mean
	// with less susceptiblity to outliers.
	var discMates []int
	for _, a := range alns {
		if a.MapQ < o.MinMappingQuality {
			continue
		}
		s := a.At(p.Pos)
		if s == nil || s.Qual < o.MinBaseQuality {
			continue
		}

		if a.Flags&sam.Paired == sam.Paired {
			if a.Flags&sam.ProperPair == sam.ProperPair {
				p.ProperPairs++
			}
			if abs(a.Start()-a.MatePos) > o.ConcordantCutoff {
				p.Discordant++
			}
			if a.MateRef.ID() != a.Ref.ID() {
				p.DiscordantChrom++
				discMates = append(discMates, a.MateRef.ID())
			} else {
				// same chromosome.
				if a.Start() < a.MatePos && a.Flags&sam.Reverse != sam.Reverse {
					p.InsertSizeLPs = append(p.InsertSizeLPs, int32(a.MatePos-a.Start()))
				} else if a.Start() > a.MatePos && a.Flags&sam.Reverse == sam.Reverse {
					p.InsertSizeRMs = append(p.InsertSizeRMs, int32(float64(a.Start()-a.MatePos)))
				} else {
					if a.Flags&sam.Reverse == sam.Reverse && a.Flags&sam.MateReverse == sam.MateReverse {
						p.OrientationMinusMinus++
					} else if 0 == a.Flags&sam.Reverse && 0 == a.Flags&sam.MateReverse {
						p.OrientationPlusPlus++
					} else {
						p.OrientationMinusPlus++
					}
				}
			}

		}
		if a.Flags&sam.Secondary == 0 {
			if tags, ok := a.Record.Tag([]byte{'S', 'A'}); ok {
				p.Splitters++
				if bytes.Count([]byte(tags), []byte{';'}) <= 1 {
					p.Splitters1++
				}

				p.updateSplitters(o, tags, a.Strand() != -1)
			}
		}

		p.Depth++

		switch s.Right.Type() {
		case sam.CigarMatch:
			break
		case sam.CigarInsertion:
			p.InsertionStarts++
		case sam.CigarSoftClipped:
			if s.Right.Len() >= o.MinClipLength {
				p.SoftStarts++
			}
		case sam.CigarHardClipped:
			if s.Right.Len() >= o.MinClipLength {
				p.HardStarts++
			}
		}

		switch s.Left.Type() {
		case sam.CigarMatch:
			break
		case sam.CigarInsertion:
			p.InsertionEnds++
		case sam.CigarSoftClipped:
			if s.Left.Len() >= o.MinClipLength {
				p.SoftEnds++
			}
		case sam.CigarHardClipped:
			if s.Left.Len() >= o.MinClipLength {
				p.HardEnds++
			}
		}

		if s.Head {
			p.Heads++
		} else if s.Tail {
			p.Tails++
		}

		if s.At.Type() == sam.CigarDeletion {
			p.Deletions++
		}
		if o.IncludeBases {
			p.Bases = append(p.Bases, s.Base)
			p.Quals = append(p.Quals, s.Qual)
		}
		if s.Base != p.RefBase {
			p.MisMatches++
		}
	}

	if p.DiscordantChrom > 1 {
		p.DiscordantChromEntropy = float32(entropy(discMates))
	}
	// don't set this if we don't know the reference base.
	if p.RefBase == 'N' {
		p.MisMatches = 0
	}
}

// track the actual positions of the splitters and check the orientation.
func (p *Pile) updateSplitters(o Options, tags []byte, readStrand bool) {
	if o.SplitterVerbosity == 0 {
		return
	}
	sas := ParseSAs(tags)
	// if there is an orientation change, we only want to count it once.
	var orientationChange bool
	for _, sa := range sas {
		if sa.MapQ >= o.MinMappingQuality {
			p.SplitterPositions = append(p.SplitterPositions, Position{Chrom: string(sa.Chrom), Start: sa.Pos, End: sa.End(), Strand: sa.Strand})
			if readStrand != sa.Strand {
				orientationChange = true
			}
		}
	}
	if orientationChange {
		p.OrientationSplitter++
	}
}

// Align is a sam.Record with a cursor to track position in the read and reference.
type Align struct {
	*sam.Record
	// track the into the Cigar Ops.
	CursorCigar int
	// track basepairs into the Reference
	CursorPos int
	// track basepairs into the read
	CursorRead int

	// Sequence holds the expanded sequence.
	Sequence []byte
	// track that we're alwasy moving forward.
	lastPos int
}

// CigarSummary is a summary of the context of a position for 1 alignment
type CigarSummary struct {
	At        sam.CigarOp
	Left      sam.CigarOp
	Right     sam.CigarOp
	Qual      uint8
	Head      bool
	Tail      bool
	Base      byte
	Insertion []byte
}

// At returns the CigarOp for a particular genomic position of the given read.
func (a *Align) At(pos0 int) *CigarSummary {
	if pos0+1 <= a.lastPos {
		log.Fatal("can't use align on the same position or lower position's that previous calls")
	}
	a.lastPos = pos0 + 1
	pos := a.Pos + a.CursorPos
	if pos0 < pos || len(a.Cigar) == 0 {
		return nil
	}

	if a.Sequence == nil {
		a.Sequence = a.Seq.Expand()
	}

	res := &CigarSummary{Left: a.Cigar[max(a.CursorCigar-1, 0)]}
	res.Head = pos0 == pos && a.CursorPos == 0

	for _, co := range a.Cigar[a.CursorCigar:] {
		t := co.Type()
		con := t.Consumes()
		l := co.Len()
		lr := l * con.Reference
		lq := l * con.Query

		// we can update the Left if the current position is still less
		// then the requested position.
		if pos0 > pos {
			res.Left = co
		}
		if pos+lr > pos0 {
			res.At = co
			right := co
			res.Right = co
			// add distance into read. Subtract how far before la we are.
			readi := a.CursorRead + (pos0 - pos)
			if lq > 0 {
				// if we consume bases from the read, then grab them.
				res.Base = a.Sequence[readi]
				res.Qual = a.Qual[readi]
			} else {
				res.Base = '*'
				if co.Type() == sam.CigarSkipped {
					res.Base = SkipBase
				}
			}
			if pos+lr-1 == pos0 {
				idx := a.CursorCigar + 1
				if idx == len(a.Cigar) {
					res.Tail = true
				} else {
					res.Right = a.Cigar[idx]
					right = a.Cigar[idx]
				}
			}
			if res.Right.Type() == sam.CigarInsertion {
				res.Insertion = a.Sequence[readi+1 : readi+1+right.Len()]
			}

			return res
		}
		res.Left = co
		a.CursorCigar++
		a.CursorPos += lr
		a.CursorRead += lq
		pos += lr
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"regexp"
	"strings"
	"time"

	mp "github.com/mackerelio/go-mackerel-plugin"
	fping "github.com/tatsushid/go-fastping"
)

type PingPlugin struct {
	Hosts    []string
	Tempfile string
}

func (pp PingPlugin) FetchMetrics() (map[string]float64, error) {
	stat := make(map[string]float64)

	pinger := fping.NewPinger()
	pinger.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
		rttMicroSec := float64(rtt.Nanoseconds()) / 1000.0 / 1000.0
		stat[escapeHostName(addr.String())] = rttMicroSec
	}

	for _, host := range pp.Hosts {
		ra, err := net.ResolveIPAddr("ip4:icmp", host)
		if err != nil {
			return nil, err
		}

		pinger.AddIPAddr(ra)
	}

	err := pinger.Run()
	if err != nil {
		return nil, err
	}

	pinger.RunLoop()

	return stat, nil
}

func (pp PingPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := []mp.Metrics{}
	for _, host := range pp.Hosts {
		metrics = append(metrics, mp.Metrics{
			Name:    escapeHostName(host),
			Label:   host,
			Diff:    false,
			Stacked: true,
		})
	}

	return map[string](mp.Graphs){
		"ping.rtt": mp.Graphs{
			Label:   "Ping Round Trip Times",
			Unit:    "float",
			Metrics: metrics,
		},
	}
}

func escapeHostName(host string) string {
	return strings.Replace(host, ".", "_", -1)
}

func validate(ipAddr string) bool {
	r := regexp.MustCompile("^\\d+\\.\\d+\\.\\d+\\.\\d+$")
	return r.MatchString(ipAddr)
}

func main() {
	optHost := flag.String("host", "127.0.0.1", "IP Address(es)")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	hosts := strings.Split(*optHost, ",")
	for _, host := range hosts {
		if !validate(host) {
			fmt.Fprintf(os.Stderr, "error: %v must be ipv4 address format\n", host)
			os.Exit(1)
		}
	}

	var pp PingPlugin
	pp.Hosts = hosts

	helper := mp.NewMackerelPlugin(pp)

	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-ping-%s", *optHost)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>Remove pinger.RunLoop()<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"regexp"
	"strings"
	"time"

	mp "github.com/mackerelio/go-mackerel-plugin"
	fping "github.com/tatsushid/go-fastping"
)

type PingPlugin struct {
	Hosts    []string
	Tempfile string
}

func (pp PingPlugin) FetchMetrics() (map[string]float64, error) {
	stat := make(map[string]float64)

	pinger := fping.NewPinger()
	pinger.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
		rttMicroSec := float64(rtt.Nanoseconds()) / 1000.0 / 1000.0
		stat[escapeHostName(addr.String())] = rttMicroSec
	}

	for _, host := range pp.Hosts {
		ra, err := net.ResolveIPAddr("ip4:icmp", host)
		if err != nil {
			return nil, err
		}

		pinger.AddIPAddr(ra)
	}

	err := pinger.Run()
	if err != nil {
		return nil, err
	}

	return stat, nil
}

func (pp PingPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := []mp.Metrics{}
	for _, host := range pp.Hosts {
		metrics = append(metrics, mp.Metrics{
			Name:    escapeHostName(host),
			Label:   host,
			Diff:    false,
			Stacked: true,
		})
	}

	return map[string](mp.Graphs){
		"ping.rtt": mp.Graphs{
			Label:   "Ping Round Trip Times",
			Unit:    "float",
			Metrics: metrics,
		},
	}
}

func escapeHostName(host string) string {
	return strings.Replace(host, ".", "_", -1)
}

func validate(ipAddr string) bool {
	r := regexp.MustCompile("^\\d+\\.\\d+\\.\\d+\\.\\d+$")
	return r.MatchString(ipAddr)
}

func main() {
	optHost := flag.String("host", "127.0.0.1", "IP Address(es)")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	hosts := strings.Split(*optHost, ",")
	for _, host := range hosts {
		if !validate(host) {
			fmt.Fprintf(os.Stderr, "error: %v must be ipv4 address format\n", host)
			os.Exit(1)
		}
	}

	var pp PingPlugin
	pp.Hosts = hosts

	helper := mp.NewMackerelPlugin(pp)

	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-ping-%s", *optHost)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>package dns

// Everything is assumed in the ClassINET class. If
// you need other classes you are on your own.

// SetReply creates a reply packet from a request message.
func (dns *Msg) SetReply(request *Msg) {
	dns.MsgHdr.Id = request.MsgHdr.Id
	dns.MsgHdr.Authoritative = true
	dns.MsgHdr.Response = true
	dns.MsgHdr.Opcode = OpcodeQuery
	dns.MsgHdr.Rcode = RcodeSuccess
	dns.Question = make([]Question, 1)
	dns.Question[0] = request.Question[0]
}

// SetQuestion creates a question packet.
func (dns *Msg) SetQuestion(z string, t uint16) {
	dns.MsgHdr.Id = Id()
	dns.MsgHdr.RecursionDesired = true
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, t, ClassINET}
}

// SetNotify creates a notify packet.
func (dns *Msg) SetNotify(z string) {
	dns.MsgHdr.Opcode = OpcodeNotify
	dns.MsgHdr.Authoritative = true
	dns.MsgHdr.Id = Id()
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, TypeSOA, ClassINET}
}

// SetRcode creates an error packet.
func (dns *Msg) SetRcode(request *Msg, rcode int) {
	dns.MsgHdr.Rcode = rcode
	dns.MsgHdr.Opcode = OpcodeQuery
	dns.MsgHdr.Response = true
	dns.MsgHdr.Authoritative = false
	dns.MsgHdr.Id = request.MsgHdr.Id
	dns.Question = make([]Question, 1)
	dns.Question[0] = request.Question[0]
}

// SetRcodeFormatError creates a packet with FormError set.
func (dns *Msg) SetRcodeFormatError(request *Msg) {
	dns.MsgHdr.Rcode = RcodeFormatError
	dns.MsgHdr.Opcode = OpcodeQuery
	dns.MsgHdr.Response = true
	dns.MsgHdr.Authoritative = false
	dns.MsgHdr.Id = request.MsgHdr.Id
}

// SetUpdate makes the message a dynamic update packet. It
// sets the ZONE section to: z, TypeSOA, classINET.
func (dns *Msg) SetUpdate(z string) {
	dns.MsgHdr.Id = Id()
	dns.MsgHdr.Opcode = OpcodeUpdate
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, TypeSOA, ClassINET}
}

// SetIxfr creates dns msg suitable for requesting an ixfr.
func (dns *Msg) SetIxfr(z string, serial uint32) {
	dns.MsgHdr.Id = Id()
	dns.Question = make([]Question, 1)
	dns.Ns = make([]RR, 1)
	s := new(RR_SOA)
	s.Hdr = RR_Header{z, TypeSOA, ClassINET, DefaultTtl, 0}
	s.Serial = serial

	dns.Question[0] = Question{z, TypeIXFR, ClassINET}
	dns.Ns[0] = s
}

// SetAxfr creates dns msg suitable for requesting an axfr.
func (dns *Msg) SetAxfr(z string) {
	dns.MsgHdr.Id = Id()
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, TypeAXFR, ClassINET}
}

// SetTsig appends a TSIG RR to the message.
// This is only a skeleton Tsig RR that is added as the last RR in the
// additional section. The caller should then call TsigGenerate,
// to generate the complete TSIG with the secret.
func (dns *Msg) SetTsig(z, algo string, fudge uint16, timesigned uint64) {
	t := new(RR_TSIG)
	t.Hdr = RR_Header{z, TypeTSIG, ClassANY, 0, 0}
	t.Algorithm = algo
	t.Fudge = 300
	t.TimeSigned = timesigned
	dns.Extra = append(dns.Extra, t)
}

// SetEdns0 appends a EDNS0 OPT RR to the message.
// TSIG should always the last RR in a message.
func (dns *Msg) SetEdns0(udpsize uint16, do bool) {
	e := new(RR_OPT)
	e.Hdr.Name = "."
	e.Hdr.Rrtype = TypeOPT
	e.SetUDPSize(udpsize)
	if do {
		e.SetDo()
	}
	dns.Extra = append(dns.Extra, e)
}

// IsRcode checks if the header of the packet has rcode set.
func (dns *Msg) IsRcode(rcode int) (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Rcode == rcode
	return
}

// IsQuestion returns true if the packet is a question.
func (dns *Msg) IsQuestion() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Response == false
	return
}

// IsRcodeFormatError checks if the message has FormErr set.
func (dns *Msg) IsRcodeFormatError() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Rcode == RcodeFormatError
	return
}

// IsUpdate checks if the message is a dynamic update packet.
func (dns *Msg) IsUpdate() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeUpdate
	ok = ok && dns.Question[0].Qtype == TypeSOA
	return
}

// IsNotify checks if the message is a valid notify packet.
func (dns *Msg) IsNotify() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeNotify
	ok = ok && dns.Question[0].Qclass == ClassINET
	ok = ok && dns.Question[0].Qtype == TypeSOA
	return
}

// IsAxfr checks if the message is a valid axfr request packet.
func (dns *Msg) IsAxfr() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeQuery
	ok = ok && dns.Question[0].Qclass == ClassINET
	ok = ok && dns.Question[0].Qtype == TypeAXFR
	return
}

// IsIXfr checks if the message is a valid ixfr request packet.
func (dns *Msg) IsIxfr() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeQuery
	ok = ok && dns.Question[0].Qclass == ClassINET
	ok = ok && dns.Question[0].Qtype == TypeIXFR
	return
}

// IsTsig checks if the message has a TSIG record as the last record
// in the additional section.
func (dns *Msg) IsTsig() (ok bool) {
	if len(dns.Extra) > 0 {
		return dns.Extra[len(dns.Extra)-1].Header().Rrtype == TypeTSIG
	}
	return
}

// IsEdns0 checks if the message has a Edns0 record, any EDNS0
// record in the additional section will do.
func (dns *Msg) IsEdns0() (ok bool) {
	for _, r := range dns.Extra {
		if r.Header().Rrtype == TypeOPT {
			return true
		}
	}
	return
}

// IsDomainName checks if s is a valid domainname, it returns
// the number of labels, total length and true, when a domain name is valid.
// When false is returned the labelcount and length are not defined.
func IsDomainName(s string) (uint8, uint8, bool) { // copied from net package.
	// See RFC 1035, RFC 3696.
	l := len(s)
	if l == 0 || l > 255 {
		return 0, 0, false
	}
	longer := 0
	// Simplify checking loop: make the name end in a dot.
	// Don't call Fqdn() to save another len(s).
	// Keep in mind that if we do this, we report a longer
	// length
	if s[l-1] != '.' {
		s += "."
		l++
		longer = 1
	}
        // Preloop check for root label
        if s == "." {
                return 0, 1, true
        }

	last := byte('.')
	ok := false // ok once we've seen a letter
	partlen := 0
	labels := uint8(0)
	for i := 0; i < l; i++ {
		c := s[i]
		switch {
		default:
			return 0, uint8(l - longer), false
		case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || c == '_' || c == '*':
			ok = true
			partlen++
		case c == '\\':
			// Ok
		case '0' <= c && c <= '9':
			// fine
			partlen++
		case c == '-':
			// byte before dash cannot be dot
			if last == '.' {
				return 0, uint8(l - longer), false
			}
			partlen++
		case c == '.':
			// byte before dot cannot be dot, dash
			if last == '.' || last == '-' {
				return 0, uint8(l - longer), false
			}
			if last == '\\' { // Ok, escaped dot.
				partlen++
				break
			}
			if partlen > 63 || partlen == 0 {
				return 0, uint8(l - longer), false
			}
			partlen = 0
			labels++
		}
		last = c
	}
	return labels, uint8(l - longer), ok
}

// IsFqdn checks if a domain name is fully qualified.
func IsFqdn(s string) bool {
	if len(s) == 0 {
		return false // ?
	}
	return s[len(s)-1] == '.'
}

// Fqdns return the fully qualified domain name from s.
// If s is already fully qualified, it behaves as the identity function.
func Fqdn(s string) string {
	if IsFqdn(s) {
		return s
	}
	return s + "."
}
<commit_msg>All numeric is OK too nowadays<commit_after>package dns

// Everything is assumed in the ClassINET class. If
// you need other classes you are on your own.

// SetReply creates a reply packet from a request message.
func (dns *Msg) SetReply(request *Msg) {
	dns.MsgHdr.Id = request.MsgHdr.Id
	dns.MsgHdr.Authoritative = true
	dns.MsgHdr.Response = true
	dns.MsgHdr.Opcode = OpcodeQuery
	dns.MsgHdr.Rcode = RcodeSuccess
	dns.Question = make([]Question, 1)
	dns.Question[0] = request.Question[0]
}

// SetQuestion creates a question packet.
func (dns *Msg) SetQuestion(z string, t uint16) {
	dns.MsgHdr.Id = Id()
	dns.MsgHdr.RecursionDesired = true
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, t, ClassINET}
}

// SetNotify creates a notify packet.
func (dns *Msg) SetNotify(z string) {
	dns.MsgHdr.Opcode = OpcodeNotify
	dns.MsgHdr.Authoritative = true
	dns.MsgHdr.Id = Id()
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, TypeSOA, ClassINET}
}

// SetRcode creates an error packet.
func (dns *Msg) SetRcode(request *Msg, rcode int) {
	dns.MsgHdr.Rcode = rcode
	dns.MsgHdr.Opcode = OpcodeQuery
	dns.MsgHdr.Response = true
	dns.MsgHdr.Authoritative = false
	dns.MsgHdr.Id = request.MsgHdr.Id
	dns.Question = make([]Question, 1)
	dns.Question[0] = request.Question[0]
}

// SetRcodeFormatError creates a packet with FormError set.
func (dns *Msg) SetRcodeFormatError(request *Msg) {
	dns.MsgHdr.Rcode = RcodeFormatError
	dns.MsgHdr.Opcode = OpcodeQuery
	dns.MsgHdr.Response = true
	dns.MsgHdr.Authoritative = false
	dns.MsgHdr.Id = request.MsgHdr.Id
}

// SetUpdate makes the message a dynamic update packet. It
// sets the ZONE section to: z, TypeSOA, classINET.
func (dns *Msg) SetUpdate(z string) {
	dns.MsgHdr.Id = Id()
	dns.MsgHdr.Opcode = OpcodeUpdate
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, TypeSOA, ClassINET}
}

// SetIxfr creates dns msg suitable for requesting an ixfr.
func (dns *Msg) SetIxfr(z string, serial uint32) {
	dns.MsgHdr.Id = Id()
	dns.Question = make([]Question, 1)
	dns.Ns = make([]RR, 1)
	s := new(RR_SOA)
	s.Hdr = RR_Header{z, TypeSOA, ClassINET, DefaultTtl, 0}
	s.Serial = serial

	dns.Question[0] = Question{z, TypeIXFR, ClassINET}
	dns.Ns[0] = s
}

// SetAxfr creates dns msg suitable for requesting an axfr.
func (dns *Msg) SetAxfr(z string) {
	dns.MsgHdr.Id = Id()
	dns.Question = make([]Question, 1)
	dns.Question[0] = Question{z, TypeAXFR, ClassINET}
}

// SetTsig appends a TSIG RR to the message.
// This is only a skeleton Tsig RR that is added as the last RR in the
// additional section. The caller should then call TsigGenerate,
// to generate the complete TSIG with the secret.
func (dns *Msg) SetTsig(z, algo string, fudge uint16, timesigned uint64) {
	t := new(RR_TSIG)
	t.Hdr = RR_Header{z, TypeTSIG, ClassANY, 0, 0}
	t.Algorithm = algo
	t.Fudge = 300
	t.TimeSigned = timesigned
	dns.Extra = append(dns.Extra, t)
}

// SetEdns0 appends a EDNS0 OPT RR to the message.
// TSIG should always the last RR in a message.
func (dns *Msg) SetEdns0(udpsize uint16, do bool) {
	e := new(RR_OPT)
	e.Hdr.Name = "."
	e.Hdr.Rrtype = TypeOPT
	e.SetUDPSize(udpsize)
	if do {
		e.SetDo()
	}
	dns.Extra = append(dns.Extra, e)
}

// IsRcode checks if the header of the packet has rcode set.
func (dns *Msg) IsRcode(rcode int) (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Rcode == rcode
	return
}

// IsQuestion returns true if the packet is a question.
func (dns *Msg) IsQuestion() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Response == false
	return
}

// IsRcodeFormatError checks if the message has FormErr set.
func (dns *Msg) IsRcodeFormatError() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Rcode == RcodeFormatError
	return
}

// IsUpdate checks if the message is a dynamic update packet.
func (dns *Msg) IsUpdate() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeUpdate
	ok = ok && dns.Question[0].Qtype == TypeSOA
	return
}

// IsNotify checks if the message is a valid notify packet.
func (dns *Msg) IsNotify() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeNotify
	ok = ok && dns.Question[0].Qclass == ClassINET
	ok = ok && dns.Question[0].Qtype == TypeSOA
	return
}

// IsAxfr checks if the message is a valid axfr request packet.
func (dns *Msg) IsAxfr() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeQuery
	ok = ok && dns.Question[0].Qclass == ClassINET
	ok = ok && dns.Question[0].Qtype == TypeAXFR
	return
}

// IsIXfr checks if the message is a valid ixfr request packet.
func (dns *Msg) IsIxfr() (ok bool) {
	if len(dns.Question) == 0 {
		return false
	}
	ok = dns.MsgHdr.Opcode == OpcodeQuery
	ok = ok && dns.Question[0].Qclass == ClassINET
	ok = ok && dns.Question[0].Qtype == TypeIXFR
	return
}

// IsTsig checks if the message has a TSIG record as the last record
// in the additional section.
func (dns *Msg) IsTsig() (ok bool) {
	if len(dns.Extra) > 0 {
		return dns.Extra[len(dns.Extra)-1].Header().Rrtype == TypeTSIG
	}
	return
}

// IsEdns0 checks if the message has a Edns0 record, any EDNS0
// record in the additional section will do.
func (dns *Msg) IsEdns0() (ok bool) {
	for _, r := range dns.Extra {
		if r.Header().Rrtype == TypeOPT {
			return true
		}
	}
	return
}

// IsDomainName checks if s is a valid domainname, it returns
// the number of labels, total length and true, when a domain name is valid.
// When false is returned the labelcount and length are not defined.
func IsDomainName(s string) (uint8, uint8, bool) { // copied from net package.
	// See RFC 1035, RFC 3696.
	l := len(s)
	if l == 0 || l > 255 {
		return 0, 0, false
	}
	longer := 0
	// Simplify checking loop: make the name end in a dot.
	// Don't call Fqdn() to save another len(s).
	// Keep in mind that if we do this, we report a longer
	// length
	if s[l-1] != '.' {
		s += "."
		l++
		longer = 1
	}
        // Preloop check for root label
        if s == "." {
                return 0, 1, true
        }

	last := byte('.')
	ok := false // ok once we've seen a letter or digit
	partlen := 0
	labels := uint8(0)
	for i := 0; i < l; i++ {
		c := s[i]
		switch {
		default:
			return 0, uint8(l - longer), false
		case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || c == '_' || c == '*':
			ok = true
			partlen++
		case c == '\\':
			// Ok
		case '0' <= c && c <= '9':
			ok = true
			partlen++
		case c == '-':
			// byte before dash cannot be dot
			if last == '.' {
				return 0, uint8(l - longer), false
			}
			partlen++
		case c == '.':
			// byte before dot cannot be dot, dash
			if last == '.' || last == '-' {
				return 0, uint8(l - longer), false
			}
			if last == '\\' { // Ok, escaped dot.
				partlen++
				break
			}
			if partlen > 63 || partlen == 0 {
				return 0, uint8(l - longer), false
			}
			partlen = 0
			labels++
		}
		last = c
	}
	return labels, uint8(l - longer), ok
}

// IsFqdn checks if a domain name is fully qualified.
func IsFqdn(s string) bool {
	if len(s) == 0 {
		return false // ?
	}
	return s[len(s)-1] == '.'
}

// Fqdns return the fully qualified domain name from s.
// If s is already fully qualified, it behaves as the identity function.
func Fqdn(s string) string {
	if IsFqdn(s) {
		return s
	}
	return s + "."
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"encoding/json"
	"hash"
	//	"fmt"
	"crypto/sha512"
	"encoding/hex"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

//Image contains data we got from API that we are using to filter and fetch said image next
type Image struct {
	Imgid          int    `json:"id_number"`
	URL            string `json:"image"`
	Filename       string
	Score          int    `json:"score"`
	OriginalFormat string `json:"original_format"`
	Hashval        string `json:"sha512_hash"`
}

//Search returns to us array of searched images...
type Search struct {
	Images []Image `json:"search"`
}

//ImageCh is a channel of image data. You can put images into channel by parsing
//Derpibooru API by id(s) or  by tags and you can download images that are already
//in channel
type ImageCh chan Image

//Push gets unmarchalled JSON info, massages it and plugs it into channel so it
//would be processed in other places
func (imgchan ImageCh) push(dat Image) {
	dat.Filename = strconv.Itoa(dat.Imgid) + "." + dat.OriginalFormat
	dat.URL = prefix + dat.URL
	if dat.OriginalFormat == "svg" {
		i := strings.LastIndex(dat.URL, ".")
		if i != -1 {
			dat.URL = dat.URL[:i] + ".svg" //Was afraid to extract things I needed from the date field, so extracting them from URL.
		}
	}
	imgchan <- dat
}

//ParseImg gets image IDs, fetches information about those images from Derpibooru and pushes them into the channel.
func (imgchan ImageCh) ParseImg() {

	for _, imgid := range opts.Args.IDs {
		source := prefix + "//derpiboo.ru/images/" + strconv.Itoa(imgid) + ".json"
		if opts.Key != "" {
			source = source + "?key=" + opts.Key
		}

		//log.Println("Getting image info at:", source)

		response, err := http.Get(source) //Getting our nice http response. Needs checking for 404 and other responses that are... less expected
		if err != nil {
			if _, ok := err.(x509.UnknownAuthorityError); ok && opts.Unsafe {
				log.Println("Warning: ", err)
			} else {
				elog.Println(err)
				continue
			}
		}

		defer func() {
			err = response.Body.Close() //and not forgetting to close it when it's done
			if err != nil {
				elog.Fatalln("Could  not close server response")
			}
		}()
		var dat Image

		body, err := ioutil.ReadAll(response.Body) //stolen from official documentation
		if err != nil {
			elog.Println(err)
			continue
		}

		if err := json.Unmarshal(body, &dat); //transforming json into native map

		err != nil {
			elog.Println(err)
			continue
		}

		imgchan.push(dat)
	}

	close(imgchan) //closing channel, we are done here

	return
}

//DlImg reads image data from channel and downloads specified images to disc
func (imgchan ImageCh) DlImg() {

	log.Println("Worker started; reading channel") //nice notification that we are not forgotten

	hasher := sha512.New() //checksums will be done in this

	for imgdata := range imgchan {

		if imgdata.Filename == "" {
			elog.Println("Empty filename. Oops?") //something somewhere had gone wrong, not a cause to die, going to the next image
			continue
		}

		log.Println("Saving as", imgdata.Filename)

		imgdata.saveImage(hasher)

		//fmt.Println("\n", hex.EncodeToString(hash.Sum(nil)), "\n", imgdata.hash )

	}
	done <- true
}

func (imgdata Image) saveImage(hasher hash.Hash) { // To not hold all the files open when there is no need. All pointers to files are in the scope of this function.

	output, err := os.Create(opts.ImageDir + string(os.PathSeparator) + imgdata.Filename) //And now, THE FILE!
	if err != err {
		elog.Println("Error when creating file for image" + strconv.Itoa(imgdata.Imgid))
		elog.Println(err) //Either we got no permisson or no space, end of line
		return
	}
	defer func() {
		err = output.Close() //Not forgetting to deal with it after completing download
		if err != nil {
			elog.Fatalln("Could  not close downloaded file")
		}
	}()

	start := time.Now()
	response, err := http.Get(imgdata.URL)
	if err != nil {
		if _, ok := err.(x509.UnknownAuthorityError); ok && opts.Unsafe {
			log.Println("Warning: ", err)
		} else {
			elog.Println("Error when getting image", imgdata.Imgid)
			elog.Println(err)
			return
		}
	}
	defer func() {
		err = response.Body.Close() //Same, we shall not listen to the void when we finished getting image
		if err != nil {
			elog.Fatalln("Could  not close server response")
		}
	}()

	size, err := io.Copy(output, io.TeeReader(response.Body, hasher)) //	Writing things we got from Derpibooru into the file and into hasher
	if err != nil {
		elog.Println("Unable to write image on disk, id ", imgdata.Imgid)
		elog.Println(err)
		return
	}
	timed := time.Since(start).Seconds()

	hash := hex.EncodeToString(hasher.Sum(nil))

	if hash != imgdata.Hashval {
		elog.Println("Hash mismatch, got ", hash, " instead of ", imgdata.Hashval)
	}

	hasher.Reset()

	log.Printf("Downloaded %d bytes in %.2fs, speed %s/s\n", size, timed, fmtbytes(float64(size)/timed))
}

//ParseTag gets image tags, fetches information about all images it could from Derpibooru and pushes them into the channel.
func (imgchan ImageCh) ParseTag() {

	source := prefix + "//derpiboo.ru/search.json?q=" + opts.Tag //yay hardwiring url strings!

	if opts.Key != "" {
		source = source + "&key=" + opts.Key
	}

	log.Println("Searching as", source)

	for i := opts.StartPage; opts.StopPage == 0 || i <= opts.StopPage; i++ {
		//I suspect that all those returns could be dealt with in some way. But lazy.
		log.Println("Searching page", i)

		response, err := http.Get(source + "&page=" + strconv.Itoa(i))
		//Getting our nice http response. Needs checking for 404 and other responses that are... less expected

		if err != nil {
			if _, ok := err.(x509.UnknownAuthorityError); ok && opts.Unsafe {
				log.Println("Warning: ", err)
			} else {
				elog.Println("Error while getting search page", i)
				elog.Println(err)
				continue
			}
		}

		defer func() {
			err = response.Body.Close() //and not forgetting to close it when it's done. And before we panic and die horribly.
			if err != nil {
				elog.Fatalln("Could  not close server response")
			}
		}()

		var dats Search //Because we got array incoming instead of single object, we using an slive of maps!

		//fmt.Println(resp)

		body, err := ioutil.ReadAll(response.Body) //stolen from official documentation
		if err != nil {
			elog.Println("Error while reading search page", i)
			elog.Println(err)
			continue
		}

		//fmt.Println(body)

		err = json.Unmarshal(body, &dats) //transforming json into native slice of maps

		if err != nil {
			elog.Println("Error while parsing search page", i)
			elog.Println(err)
			continue

		}

		if len(dats.Images) == 0 {
			log.Println("Pages are over") //Does not mean that process is over.
			break
		} //exit due to finishing all pages

		for _, dat := range dats.Images {
			imgchan.push(dat)
		}

	}

	close(imgchan)
}
<commit_msg>some additional error checks<commit_after>package main

import (
	"crypto/x509"
	"encoding/json"
	"hash"
	//	"fmt"
	"crypto/sha512"
	"encoding/hex"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

//Image contains data we got from API that we are using to filter and fetch said image next
type Image struct {
	Imgid          int    `json:"id_number"`
	URL            string `json:"image"`
	Filename       string
	Score          int    `json:"score"`
	OriginalFormat string `json:"original_format"`
	Hashval        string `json:"sha512_hash"`
}

//Search returns to us array of searched images...
type Search struct {
	Images []Image `json:"search"`
}

//ImageCh is a channel of image data. You can put images into channel by parsing
//Derpibooru API by id(s) or  by tags and you can download images that are already
//in channel
type ImageCh chan Image

//Push gets unmarchalled JSON info, massages it and plugs it into channel so it
//would be processed in other places
func (imgchan ImageCh) push(dat Image) {
	dat.Filename = strconv.Itoa(dat.Imgid) + "." + dat.OriginalFormat
	dat.URL = prefix + dat.URL
	if dat.OriginalFormat == "svg" {
		i := strings.LastIndex(dat.URL, ".")
		if i != -1 {
			dat.URL = dat.URL[:i] + ".svg" //Was afraid to extract things I needed from the date field, so extracting them from URL.
		}
	}
	imgchan <- dat
}

//ParseImg gets image IDs, fetches information about those images from Derpibooru and pushes them into the channel.
func (imgchan ImageCh) ParseImg() {

	for _, imgid := range opts.Args.IDs {
		source := prefix + "//derpiboo.ru/images/" + strconv.Itoa(imgid) + ".json"
		if opts.Key != "" {
			source = source + "?key=" + opts.Key
		}

		//log.Println("Getting image info at:", source)

		response, err := http.Get(source) //Getting our nice http response. Needs checking for 404 and other responses that are... less expected
		if err != nil {
			if _, ok := err.(x509.UnknownAuthorityError); ok && opts.Unsafe {
				log.Println("Warning: ", err)
			} else {
				elog.Println(err)
				continue
			}
		}

		defer func() {
			err = response.Body.Close() //and not forgetting to close it when it's done
			if err != nil {
				elog.Fatalln("Could  not close server response")
			}
		}()
		var dat Image

		body, err := ioutil.ReadAll(response.Body) //stolen from official documentation
		if err != nil {
			elog.Println(err)
			continue
		}

		if err := json.Unmarshal(body, &dat); //transforming json into native map

		err != nil {
			elog.Println(err)
			continue
		}

		imgchan.push(dat)
	}

	close(imgchan) //closing channel, we are done here

	return
}

//DlImg reads image data from channel and downloads specified images to disc
func (imgchan ImageCh) DlImg() {

	log.Println("Worker started; reading channel") //nice notification that we are not forgotten

	hasher := sha512.New() //checksums will be done in this

	for imgdata := range imgchan {

		if imgdata.Filename == "" {
			elog.Println("Empty filename. Oops?") //something somewhere had gone wrong, not a cause to die, going to the next image
			continue
		}

		log.Println("Saving as", imgdata.Filename)

		imgdata.saveImage(hasher)

		//fmt.Println("\n", hex.EncodeToString(hash.Sum(nil)), "\n", imgdata.hash )

	}
	done <- true
}

func (imgdata Image) saveImage(hasher hash.Hash) { // To not hold all the files open when there is no need. All pointers to files are in the scope of this function.

	output, err := os.Create(opts.ImageDir + string(os.PathSeparator) + imgdata.Filename) //And now, THE FILE!
	if err != err {
		elog.Println("Error when creating file for image" + strconv.Itoa(imgdata.Imgid))
		elog.Println(err) //Either we got no permisson or no space, end of line
		return
	}
	defer func() {
		err = output.Close() //Not forgetting to deal with it after completing download
		if err != nil {
			elog.Fatalln("Could  not close downloaded file")
		}
	}()

	start := time.Now()
	response, err := http.Get(imgdata.URL)
	if err != nil {
		if _, ok := err.(x509.UnknownAuthorityError); ok && opts.Unsafe {
			log.Println("Warning: ", err)
		} else {
			elog.Println("Error when getting image", imgdata.Imgid)
			elog.Println(err)
			return
		}
	}
	defer func() {
		err = response.Body.Close() //Same, we shall not listen to the void when we finished getting image
		if err != nil {
			elog.Fatalln("Could  not close server response")
		}
	}()

	size, err := io.Copy(output, io.TeeReader(response.Body, hasher)) //	Writing things we got from Derpibooru into the file and into hasher
	if err != nil {
		elog.Println("Unable to write image on disk, id ", imgdata.Imgid)
		elog.Println(err)
		return
	}
	timed := time.Since(start).Seconds()

	hash := hex.EncodeToString(hasher.Sum(nil))

	if hash != imgdata.Hashval {
		elog.Println("Hash mismatch, got ", hash, " instead of ", imgdata.Hashval)
	}

	hasher.Reset()

	log.Printf("Downloaded %d bytes in %.2fs, speed %s/s\n", size, timed, fmtbytes(float64(size)/timed))
}

//ParseTag gets image tags, fetches information about all images it could from Derpibooru and pushes them into the channel.
func (imgchan ImageCh) ParseTag() {

	source := prefix + "//derpiboo.ru/search.json?q=" + opts.Tag //yay hardwiring url strings!

	if opts.Key != "" {
		source = source + "&key=" + opts.Key
	}

	log.Println("Searching as", source)

	for i := opts.StartPage; opts.StopPage == 0 || i <= opts.StopPage; i++ {
		//I suspect that all those returns could be dealt with in some way. But lazy.
		log.Println("Searching page", i)

		response, err := http.Get(source + "&page=" + strconv.Itoa(i))
		//Getting our nice http response. Needs checking for 404 and other responses that are... less expected

		if err != nil {
			if _, ok := err.(x509.UnknownAuthorityError); ok && opts.Unsafe {
				log.Println("Warning: ", err)
			} else {
				elog.Println("Error while getting search page", i)
				elog.Println(err)
				continue
			}
		}

		defer func() {
			err = response.Body.Close() //and not forgetting to close it when it's done. And before we panic and die horribly.
			if err != nil {
				elog.Fatalln("Could  not close server response")
			}
		}()

		var dats Search //Because we got array incoming instead of single object, we using an slive of maps!

		//fmt.Println(resp)

		body, err := ioutil.ReadAll(response.Body) //stolen from official documentation
		if err != nil {
			elog.Println("Error while reading search page", i)
			elog.Println(err)
			continue
		}

		//fmt.Println(body)

		err = json.Unmarshal(body, &dats) //transforming json into native slice of maps

		if err != nil {
			elog.Println("Error while parsing search page", i)
			elog.Println(err)
			if serr, ok := err.(*json.SyntaxError); ok {
				log.Println("Occurred at offset:", serr.Offset)
			}
			continue

		}

		if len(dats.Images) == 0 {
			log.Println("Pages are over") //Does not mean that process is over.
			break
		} //exit due to finishing all pages

		for _, dat := range dats.Images {
			imgchan.push(dat)
		}

	}

	close(imgchan)
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"syscall"
	"time"
)

const (
	// Service defaults.
	DefaultStopSignal  = syscall.SIGINT
	DefaultStopTimeout = 5 * time.Second
	DefaultRestart     = true
	DefaultRetries     = 3

	// Service commands.
	Start    = "start"
	Stop     = "stop"
	Restart  = "restart"
	Shutdown = "shutdown"

	// Service states.
	Starting = "starting"
	Running  = "running"
	Stopping = "stopping"
	Stopped  = "stopped"
	Exited   = "exited"
	//TODO: Implement Backoff state.
)

// Command is sent to a Service to initiate a state change.
type Command struct {
	Name     string
	Response chan<- Response
}

// respond creates and sends a command Response.
func (cmd Command) respond(service *Service, err error) {
	if cmd.Response != nil {
		cmd.Response <- Response{service, cmd.Name, err}
	}
}

// Response contains the result of a Command.
type Response struct {
	Service *Service
	Name    string
	Error   error
}

// Success returns True if the Command was successful.
func (r Response) Success() bool {
	return r.Error == nil
}

// Event is sent by a Service on a state change.
type Event struct {
	Service *Service
	State   string
}

// Service represents a controllable process. Exported fields may be set to configure the service.
type Service struct {
	Directory   string         // The process's working directory. Defaults to the current directory.
	Environment []string       // The environment of the process. Defaults to nil which indicatesA the current environment.
	StopSignal  syscall.Signal // The signal to send when stopping the process. Defaults to SIGINT.
	StopTimeout time.Duration  // How long to wait for a process to stop before sending a SIGKILL. Defaults to 5s.
	Restart     bool           // Whether or not to restart the process if it exits unexpectedly. Defaults to true.
	Retries     int            // How many times to restart a process if it fails to start. Defaults to 3.
	Stdout      io.Writer      // Where to send the process's stdout. Defaults to /dev/null.
	Stderr      io.Writer      // Where to send the process's stderr. Defaults to /dev/null.
	args        []string       // The command line of the process to run.
	command     *exec.Cmd      // The os/exec command running the process.
	state       string         // The state of the Service.
}

// New creates a new service with the default configution.
func NewService(args []string) (svc *Service, err error) {
	if cwd, err := os.Getwd(); err == nil {
		svc = &Service{
			cwd,
			nil,
			DefaultStopSignal,
			DefaultStopTimeout,
			DefaultRestart,
			DefaultRetries,
			nil,
			nil,
			args,
			nil,
			Stopped,
		}
	}
	return
}

// State gets the current state of the service.
func (s Service) State() string {
	return s.state
}

// Pid gets the PID of the service or 0 if not Running or Stopping.
func (s Service) Pid() int {
	if s.state != Running && s.state != Stopping {
		return 0
	}
	return s.command.Process.Pid
}

func (s Service) makeCommand() *exec.Cmd {
	cmd := exec.Command(s.args[0], s.args[1:]...)
	cmd.Stdout = s.Stdout
	cmd.Stderr = s.Stderr
	cmd.Stdin = nil
	cmd.Env = s.Environment
	cmd.Dir = s.Directory
	return cmd
}

func (s *Service) startProcess(states chan string) (err error) {
	defer func() {
		if err != nil {
			//TODO: Do something with this error.
			states <- Exited
		}
	}()

	s.command = s.makeCommand()
	if err = s.command.Start(); err != nil {
		return
	}

	states <- Running
	s.command.Wait()
	states <- Exited
	return
}

func (s *Service) Run(commands <-chan Command, events chan<- Event) {
	var lastCommand *Command
	states := make(chan string)
	quit := make(chan bool, 2)
	kill := make(chan int, 2)
	retries := 0

	defer func() {
		close(states)
		close(quit)
		close(kill)
	}()

	sendEvent := func(state string) {
		s.state = state
		events <- Event{s, state}
	}

	sendInvalidCmd := func(cmd *Command, state string) {
		if cmd != nil {
			cmd.respond(s, errors.New(fmt.Sprintf("invalid state transition: %s -> %s", s.state, state)))
		}
	}

	start := func(cmd *Command) {
		if s.state != Stopped && s.state != Exited {
			sendInvalidCmd(cmd, Starting)
			return
		}

		sendEvent(Starting)
		go s.startProcess(states)
	}

	stop := func(cmd *Command) {
		if s.state != Running {
			sendInvalidCmd(cmd, Stopping)
			return
		}

		sendEvent(Stopping)
		pid := s.Pid()
		s.command.Process.Signal(s.StopSignal) //TODO: Check for error.
		go func() {
			time.Sleep(s.StopTimeout)
			defer func() {
				if err := recover(); err != nil {
					if _, ok := err.(runtime.Error); !ok {
						panic(err)
					}
				}
			}()
			kill <- pid
		}()
	}

	shutdown := func(cmd *Command, lastCmd *Command) {
		if lastCmd != nil {
			lastCmd.respond(s, errors.New("service is shutting down"))
		}
		if s.state == Stopped || s.state == Exited {
			quit <- true
		} else if s.state == Running {
			stop(cmd)
		}
	}

	onRunning := func(cmd *Command) {
		sendEvent(Running)
		if cmd != nil {
			switch cmd.Name {
			case Start:
				fallthrough
			case Restart:
				cmd.respond(s, nil)
			case Shutdown:
				stop(cmd)
			}
		}
		retries = 0
	}

	onStopped := func(cmd *Command) {
		sendEvent(Stopped)
		if cmd != nil {
			switch cmd.Name {
			case Restart:
				start(cmd)
			case Stop:
				cmd.respond(s, nil)
			case Shutdown:
				quit <- true
			}
		}
	}

	onExited := func(cmd *Command, retries int) bool {
		sendEvent(Exited)
		if s.Restart && retries < s.Retries {
			start(cmd)
			return true
		}
		return false
	}

loop:
	for {
		select {
		case state := <-states:
			// running, exited
			switch state {
			case Running:
				onRunning(lastCommand)
			case Exited:
				if s.state == Stopping {
					onStopped(lastCommand)
				} else {
					if onExited(lastCommand, retries) {
						retries++
					}
				}
			}
			if lastCommand != nil {
				if lastCommand.Name == Restart && s.state == Running {
					lastCommand = nil
				} else if lastCommand.Name != Restart && lastCommand.Name != Shutdown {
					lastCommand = nil
				}
			}
		case command := <-commands:
			if lastCommand == nil || lastCommand.Name != Shutdown { // Shutdown cannot be overriden!
				switch command.Name {
				case Start:
					start(&command)
				case Stop:
					stop(&command)
				case Restart:
					stop(&command)
				case Shutdown:
					shutdown(&command, lastCommand)
				}
				lastCommand = &command
			} else {
				command.respond(s, errors.New("service is shutting down"))
			}
		case <-quit:
			if lastCommand != nil {
				lastCommand.respond(s, nil)
			}
			break loop
		case pid := <-kill:
			if pid == s.Pid() {
				s.command.Process.Kill() //TODO: Check for error.
			}
		}
	}
}
<commit_msg>Implement Backoff state and modify restart behavior.<commit_after>package service

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"syscall"
	"time"
)

const (
	// Service defaults.
	DefaultStartTimeout = 1 * time.Second
	DefaultStartRetries = 3
	DefaultStopSignal   = syscall.SIGINT
	DefaultStopTimeout  = 5 * time.Second
	DefaultStopRestart  = true

	// Service commands.
	Start    = "start"
	Stop     = "stop"
	Restart  = "restart"
	Shutdown = "shutdown"

	// Service states.
	Starting = "starting"
	Running  = "running"
	Stopping = "stopping"
	Stopped  = "stopped"
	Exited   = "exited"
	Backoff  = "backoff"
)

// Command is sent to a Service to initiate a state change.
type Command struct {
	Name     string
	Response chan<- Response
}

// respond creates and sends a command Response.
func (cmd Command) respond(service *Service, err error) {
	if cmd.Response != nil {
		cmd.Response <- Response{service, cmd.Name, err}
	}
}

// Response contains the result of a Command.
type Response struct {
	Service *Service
	Name    string
	Error   error
}

// Success returns True if the Command was successful.
func (r Response) Success() bool {
	return r.Error == nil
}

// Event is sent by a Service on a state change.
type Event struct {
	Service *Service
	State   string
}

// Service represents a controllable process. Exported fields may be set to configure the service.
type Service struct {
	Directory    string         // The process's working directory. Defaults to the current directory.
	Environment  []string       // The environment of the process. Defaults to nil which indicatesA the current environment.
	StartTimeout time.Duration  // How long the process has to run before it's considered Running.
	StartRetries int            // How many times to restart a process if it fails to start. Defaults to 3.
	StopSignal   syscall.Signal // The signal to send when stopping the process. Defaults to SIGINT.
	StopTimeout  time.Duration  // How long to wait for a process to stop before sending a SIGKILL. Defaults to 5s.
	StopRestart  bool           // Whether or not to restart the process if it exits unexpectedly. Defaults to true.
	Stdout       io.Writer      // Where to send the process's stdout. Defaults to /dev/null.
	Stderr       io.Writer      // Where to send the process's stderr. Defaults to /dev/null.
	args         []string       // The command line of the process to run.
	command      *exec.Cmd      // The os/exec command running the process.
	state        string         // The state of the Service.
}

// New creates a new service with the default configution.
func NewService(args []string) (svc *Service, err error) {
	if cwd, err := os.Getwd(); err == nil {
		svc = &Service{
			cwd,
			nil,
			DefaultStartTimeout,
			DefaultStartRetries,
			DefaultStopSignal,
			DefaultStopTimeout,
			DefaultStopRestart,
			nil,
			nil,
			args,
			nil,
			Stopped,
		}
	}
	return
}

// State gets the current state of the service.
func (s Service) State() string {
	return s.state
}

// Pid gets the PID of the service or 0 if not Running or Stopping.
func (s Service) Pid() int {
	if s.state != Running && s.state != Stopping {
		return 0
	}
	return s.command.Process.Pid
}

func (s Service) makeCommand() *exec.Cmd {
	cmd := exec.Command(s.args[0], s.args[1:]...)
	cmd.Stdout = s.Stdout
	cmd.Stderr = s.Stderr
	cmd.Stdin = nil
	cmd.Env = s.Environment
	cmd.Dir = s.Directory
	return cmd
}

func (s *Service) Run(commands <-chan Command, events chan<- Event) {
	var lastCommand *Command
	states := make(chan string)
	quit := make(chan bool, 2)
	kill := make(chan int, 2)
	retries := 0

	defer func() {
		close(states)
		close(quit)
		close(kill)
	}()

	sendEvent := func(state string) {
		s.state = state
		events <- Event{s, state}
	}

	sendInvalidCmd := func(cmd *Command, state string) {
		if cmd != nil {
			cmd.respond(s, errors.New(fmt.Sprintf("invalid state transition: %s -> %s", s.state, state)))
		}
	}

	start := func(cmd *Command) {
		if s.state != Stopped && s.state != Exited && s.state != Backoff {
			sendInvalidCmd(cmd, Starting)
			return
		}

		sendEvent(Starting)
		go func() {
			s.command = s.makeCommand()
			startTime := time.Now()
			if err := s.command.Start(); err == nil { //TODO: Don't swallow this error.
				states <- Running
				s.command.Wait()
				if time.Now().Sub(startTime) > s.StartTimeout {
					states <- Backoff
					return
				}
			}
			states <- Exited
		}()
	}

	stop := func(cmd *Command) {
		if s.state != Running {
			sendInvalidCmd(cmd, Stopping)
			return
		}

		sendEvent(Stopping)
		pid := s.Pid()
		s.command.Process.Signal(s.StopSignal) //TODO: Check for error.
		go func() {
			time.Sleep(s.StopTimeout)
			defer func() {
				if err := recover(); err != nil {
					if _, ok := err.(runtime.Error); !ok {
						panic(err)
					}
				}
			}()
			kill <- pid
		}()
	}

	shutdown := func(cmd *Command, lastCmd *Command) {
		if lastCmd != nil {
			lastCmd.respond(s, errors.New("service is shutting down"))
		}
		if s.state == Stopped || s.state == Exited {
			quit <- true
		} else if s.state == Running {
			stop(cmd)
		}
	}

	onRunning := func(cmd *Command) {
		sendEvent(Running)
		if cmd != nil {
			switch cmd.Name {
			case Start:
				fallthrough
			case Restart:
				cmd.respond(s, nil)
			case Shutdown:
				stop(cmd)
			}
		}
	}

	onStopped := func(cmd *Command) {
		sendEvent(Stopped)
		if cmd != nil {
			switch cmd.Name {
			case Restart:
				start(cmd)
			case Stop:
				cmd.respond(s, nil)
			case Shutdown:
				quit <- true
			}
		}
	}

	onExited := func(cmd *Command) {
		sendEvent(Exited)
		if s.StopRestart {
			start(cmd)
		}
	}

	onBackoff := func(cmd *Command) {
		if retries < s.StartRetries {
			sendEvent(Backoff)
			start(cmd)
			retries++
		} else {
			sendEvent(Exited)
			retries = 0
		}
	}

loop:
	for {
		select {
		case state := <-states:
			// running, exited
			switch state {
			case Running:
				onRunning(lastCommand)
			case Exited:
				if s.state == Stopping {
					onStopped(lastCommand)
				} else {
					onExited(lastCommand)
				}
			case Backoff:
				onBackoff(lastCommand)
			}
			if lastCommand != nil {
				if lastCommand.Name == Restart && s.state == Running {
					lastCommand = nil
				} else if lastCommand.Name != Restart && lastCommand.Name != Shutdown {
					lastCommand = nil
				}
			}
		case command := <-commands:
			if lastCommand == nil || lastCommand.Name != Shutdown { // Shutdown cannot be overriden!
				switch command.Name {
				case Start:
					start(&command)
				case Stop:
					stop(&command)
				case Restart:
					stop(&command)
				case Shutdown:
					shutdown(&command, lastCommand)
				}
				lastCommand = &command
			} else {
				command.respond(s, errors.New("service is shutting down"))
			}
		case <-quit:
			if lastCommand != nil {
				lastCommand.respond(s, nil)
			}
			break loop
		case pid := <-kill:
			if pid == s.Pid() {
				s.command.Process.Kill() //TODO: Check for error.
			}
		}
	}
}
<|endoftext|>