<commit_before>package server

import (
	"context"
	"fmt"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-log"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/transport"

	"github.com/micro/util/go/lib/addr"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
	// used for first registration
	registered bool
	// graceful exit
	wg sync.WaitGroup
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		// close socket
		sock.Close()

		if r := recover(); r != nil {
			log.Log("panic recovered: ", r)
			log.Log(string(debug.Stack()))
		}
	}()

	for {
		var msg transport.Message
		if err := sock.Recv(&msg); err != nil {
			return
		}

		// add to wait group
		s.wg.Add(1)

		// we use this Timeout header to set a server deadline
		to := msg.Header["Timeout"]
		// we use this Content-Type header to identify the codec needed
		ct := msg.Header["Content-Type"]

		cf, err := s.newCodec(ct)
		// TODO: needs better error handling
		if err != nil {
			sock.Send(&transport.Message{
				Header: map[string]string{
					"Content-Type": "text/plain",
				},
				Body: []byte(err.Error()),
			})
			s.wg.Done()
			return
		}

		codec := newRpcPlusCodec(&msg, sock, cf)

		// strip our headers
		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		delete(hdr, "Timeout")

		ctx := metadata.NewContext(context.Background(), hdr)

		// set the timeout if we have it
		if len(to) > 0 {
			if n, err := strconv.ParseUint(to, 10, 64); err == nil {
				ctx, _ = context.WithTimeout(ctx, time.Duration(n))
			}
		}

		// TODO: needs better error handling
		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
			s.wg.Done()
			log.Logf("Unexpected error serving request, closing socket: %v", err)
			return
		}
		s.wg.Done()
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	// update internal server
	s.rpc = &server{
		name:         s.opts.Name,
		serviceMap:   s.rpc.serviceMap,
		hdlrWrappers: s.opts.HdlrWrappers,
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}, opts ...HandlerOption) Handler {
	return newRpcHandler(h, opts...)
}

func (s *rpcServer) Handle(h Handler) error {
	s.Lock()
	defer s.Unlock()

	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}

	s.handlers[h.Name()] = h

	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}, opts ...SubscriberOption) Subscriber {
	return newSubscriber(topic, sb, opts...)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()
	node.Metadata["registry"] = config.Registry.String()

	s.RLock()
	// Maps are ordered randomly, sort the keys for consistency
	var handlerList []string
	for n, e := range s.handlers {
		// Only advertise non internal handlers
		if !e.Options().Internal {
			handlerList = append(handlerList, n)
		}
	}
	sort.Strings(handlerList)

	var subscriberList []*subscriber
	for e := range s.subscribers {
		// Only advertise non internal subscribers
		if !e.Options().Internal {
			subscriberList = append(subscriberList, e)
		}
	}
	sort.Slice(subscriberList, func(i, j int) bool {
		return subscriberList[i].topic > subscriberList[j].topic
	})

	var endpoints []*registry.Endpoint
	for _, n := range handlerList {
		endpoints = append(endpoints, s.handlers[n].Endpoints()...)
	}
	for _, e := range subscriberList {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	s.Lock()
	registered := s.registered
	s.Unlock()

	if !registered {
		log.Logf("Registering node: %s", node.Id)
	}

	// create registry options
	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}

	if err := config.Registry.Register(service, rOpts...); err != nil {
		return err
	}

	// already registered? don't need to register subscribers
	if registered {
		return nil
	}

	s.Lock()
	defer s.Unlock()

	s.registered = true

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		var opts []broker.SubscribeOption
		if queue := sb.Options().Queue; len(queue) > 0 {
			opts = append(opts, broker.Queue(queue))
		}
		sub, err := config.Broker.Subscribe(sb.Topic(), handler, opts...)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Logf("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()

	if !s.registered {
		s.Unlock()
		return nil
	}

	s.registered = false

	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Logf("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}

	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerDebugHandler(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Logf("Listening on %s", ts.Addr())
	s.Lock()
	// swap address
	addr := s.opts.Address
	s.opts.Address = ts.Addr()
	s.Unlock()

	go ts.Accept(s.accept)

	go func() {
		// wait for exit
		ch := <-s.exit

		// wait for requests to finish
		if wait(s.opts.Context) {
			s.wg.Wait()
		}

		// close transport listener
		ch <- ts.Close()

		// disconnect the broker
		config.Broker.Disconnect()

		s.Lock()
		// swap back address
		s.opts.Address = addr
		s.Unlock()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<commit_msg>Add Local/Remote ip to metadata<commit_after>package server

import (
	"context"
	"fmt"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-log"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/transport"

	"github.com/micro/util/go/lib/addr"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
	// used for first registration
	registered bool
	// graceful exit
	wg sync.WaitGroup
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		// close socket
		sock.Close()

		if r := recover(); r != nil {
			log.Log("panic recovered: ", r)
			log.Log(string(debug.Stack()))
		}
	}()

	for {
		var msg transport.Message
		if err := sock.Recv(&msg); err != nil {
			return
		}

		// add to wait group
		s.wg.Add(1)

		// we use this Timeout header to set a server deadline
		to := msg.Header["Timeout"]
		// we use this Content-Type header to identify the codec needed
		ct := msg.Header["Content-Type"]

		cf, err := s.newCodec(ct)
		// TODO: needs better error handling
		if err != nil {
			sock.Send(&transport.Message{
				Header: map[string]string{
					"Content-Type": "text/plain",
				},
				Body: []byte(err.Error()),
			})
			s.wg.Done()
			return
		}

		codec := newRpcPlusCodec(&msg, sock, cf)

		// strip our headers
		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		delete(hdr, "Timeout")

		// set local/remote ips
		hdr["Local"] = sock.Local()
		hdr["Remote"] = sock.Remote()

		ctx := metadata.NewContext(context.Background(), hdr)

		// set the timeout if we have it
		if len(to) > 0 {
			if n, err := strconv.ParseUint(to, 10, 64); err == nil {
				ctx, _ = context.WithTimeout(ctx, time.Duration(n))
			}
		}

		// TODO: needs better error handling
		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
			s.wg.Done()
			log.Logf("Unexpected error serving request, closing socket: %v", err)
			return
		}
		s.wg.Done()
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	// update internal server
	s.rpc = &server{
		name:         s.opts.Name,
		serviceMap:   s.rpc.serviceMap,
		hdlrWrappers: s.opts.HdlrWrappers,
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}, opts ...HandlerOption) Handler {
	return newRpcHandler(h, opts...)
}

func (s *rpcServer) Handle(h Handler) error {
	s.Lock()
	defer s.Unlock()

	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}

	s.handlers[h.Name()] = h

	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}, opts ...SubscriberOption) Subscriber {
	return newSubscriber(topic, sb, opts...)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()
	node.Metadata["registry"] = config.Registry.String()

	s.RLock()
	// Maps are ordered randomly, sort the keys for consistency
	var handlerList []string
	for n, e := range s.handlers {
		// Only advertise non internal handlers
		if !e.Options().Internal {
			handlerList = append(handlerList, n)
		}
	}
	sort.Strings(handlerList)

	var subscriberList []*subscriber
	for e := range s.subscribers {
		// Only advertise non internal subscribers
		if !e.Options().Internal {
			subscriberList = append(subscriberList, e)
		}
	}
	sort.Slice(subscriberList, func(i, j int) bool {
		return subscriberList[i].topic > subscriberList[j].topic
	})

	var endpoints []*registry.Endpoint
	for _, n := range handlerList {
		endpoints = append(endpoints, s.handlers[n].Endpoints()...)
	}
	for _, e := range subscriberList {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	s.Lock()
	registered := s.registered
	s.Unlock()

	if !registered {
		log.Logf("Registering node: %s", node.Id)
	}

	// create registry options
	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}

	if err := config.Registry.Register(service, rOpts...); err != nil {
		return err
	}

	// already registered? don't need to register subscribers
	if registered {
		return nil
	}

	s.Lock()
	defer s.Unlock()

	s.registered = true

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		var opts []broker.SubscribeOption
		if queue := sb.Options().Queue; len(queue) > 0 {
			opts = append(opts, broker.Queue(queue))
		}
		sub, err := config.Broker.Subscribe(sb.Topic(), handler, opts...)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Logf("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()

	if !s.registered {
		s.Unlock()
		return nil
	}

	s.registered = false

	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Logf("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}

	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerDebugHandler(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Logf("Listening on %s", ts.Addr())
	s.Lock()
	// swap address
	addr := s.opts.Address
	s.opts.Address = ts.Addr()
	s.Unlock()

	go ts.Accept(s.accept)

	go func() {
		// wait for exit
		ch := <-s.exit

		// wait for requests to finish
		if wait(s.opts.Context) {
			s.wg.Wait()
		}

		// close transport listener
		ch <- ts.Close()

		// disconnect the broker
		config.Broker.Disconnect()

		s.Lock()
		// swap back address
		s.opts.Address = addr
		s.Unlock()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"fmt"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-log"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/transport"

	"github.com/micro/util/go/lib/addr"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
	// used for first registration
	registered bool
	// graceful exit
	wg sync.WaitGroup
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		// close socket
		sock.Close()

		if r := recover(); r != nil {
			log.Log("panic recovered: ", r)
			log.Log(string(debug.Stack()))
		}
	}()

	for {
		var msg transport.Message
		if err := sock.Recv(&msg); err != nil {
			return
		}

		// add to wait group
		s.wg.Add(1)

		// we use this Timeout header to set a server deadline
		to := msg.Header["Timeout"]
		// we use this Content-Type header to identify the codec needed
		ct := msg.Header["Content-Type"]

		cf, err := s.newCodec(ct)
		// TODO: needs better error handling
		if err != nil {
			sock.Send(&transport.Message{
				Header: map[string]string{
					"Content-Type": "text/plain",
				},
				Body: []byte(err.Error()),
			})
			s.wg.Done()
			return
		}

		codec := newRpcPlusCodec(&msg, sock, cf)

		// strip our headers
		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		delete(hdr, "Timeout")

		// set local/remote ips
		hdr["Local"] = sock.Local()
		hdr["Remote"] = sock.Remote()

		ctx := metadata.NewContext(context.Background(), hdr)

		// set the timeout if we have it
		if len(to) > 0 {
			if n, err := strconv.ParseUint(to, 10, 64); err == nil {
				ctx, _ = context.WithTimeout(ctx, time.Duration(n))
			}
		}

		// TODO: needs better error handling
		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
			s.wg.Done()
			log.Logf("Unexpected error serving request, closing socket: %v", err)
			return
		}
		s.wg.Done()
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	// update internal server
	s.rpc = &server{
		name:         s.opts.Name,
		serviceMap:   s.rpc.serviceMap,
		hdlrWrappers: s.opts.HdlrWrappers,
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}, opts ...HandlerOption) Handler {
	return newRpcHandler(h, opts...)
}

func (s *rpcServer) Handle(h Handler) error {
	s.Lock()
	defer s.Unlock()

	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}

	s.handlers[h.Name()] = h

	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}, opts ...SubscriberOption) Subscriber {
	return newSubscriber(topic, sb, opts...)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()
	node.Metadata["registry"] = config.Registry.String()

	s.RLock()
	// Maps are ordered randomly, sort the keys for consistency
	var handlerList []string
	for n, e := range s.handlers {
		// Only advertise non internal handlers
		if !e.Options().Internal {
			handlerList = append(handlerList, n)
		}
	}
	sort.Strings(handlerList)

	var subscriberList []*subscriber
	for e := range s.subscribers {
		// Only advertise non internal subscribers
		if !e.Options().Internal {
			subscriberList = append(subscriberList, e)
		}
	}
	sort.Slice(subscriberList, func(i, j int) bool {
		return subscriberList[i].topic > subscriberList[j].topic
	})

	var endpoints []*registry.Endpoint
	for _, n := range handlerList {
		endpoints = append(endpoints, s.handlers[n].Endpoints()...)
	}
	for _, e := range subscriberList {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	s.Lock()
	registered := s.registered
	s.Unlock()

	if !registered {
		log.Logf("Registering node: %s", node.Id)
	}

	// create registry options
	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}

	if err := config.Registry.Register(service, rOpts...); err != nil {
		return err
	}

	// already registered? don't need to register subscribers
	if registered {
		return nil
	}

	s.Lock()
	defer s.Unlock()

	s.registered = true

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		var opts []broker.SubscribeOption
		if queue := sb.Options().Queue; len(queue) > 0 {
			opts = append(opts, broker.Queue(queue))
		}
		sub, err := config.Broker.Subscribe(sb.Topic(), handler, opts...)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Logf("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()

	if !s.registered {
		s.Unlock()
		return nil
	}

	s.registered = false

	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Logf("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}

	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerDebugHandler(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Logf("Listening on %s", ts.Addr())
	s.Lock()
	// swap address
	addr := s.opts.Address
	s.opts.Address = ts.Addr()
	s.Unlock()

	exit := make(chan bool, 1)

	go func() {
		for {
			err := ts.Accept(s.accept)

			// check if we're supposed to exit
			select {
			case <-exit:
				return
			default:
			}

			// check the error and backoff
			if err != nil {
				log.Logf("Accept error: %v", err)
				time.Sleep(time.Second)
			}

			// no error just exit
			return
		}
	}()

	go func() {
		// wait for exit
		ch := <-s.exit
		exit <- true

		// wait for requests to finish
		if wait(s.opts.Context) {
			s.wg.Wait()
		}

		// close transport listener
		ch <- ts.Close()

		// disconnect the broker
		config.Broker.Disconnect()

		s.Lock()
		// swap back address
		s.opts.Address = addr
		s.Unlock()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<commit_msg>add continue<commit_after>package server

import (
	"context"
	"fmt"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-log"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/transport"

	"github.com/micro/util/go/lib/addr"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
	// used for first registration
	registered bool
	// graceful exit
	wg sync.WaitGroup
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		// close socket
		sock.Close()

		if r := recover(); r != nil {
			log.Log("panic recovered: ", r)
			log.Log(string(debug.Stack()))
		}
	}()

	for {
		var msg transport.Message
		if err := sock.Recv(&msg); err != nil {
			return
		}

		// add to wait group
		s.wg.Add(1)

		// we use this Timeout header to set a server deadline
		to := msg.Header["Timeout"]
		// we use this Content-Type header to identify the codec needed
		ct := msg.Header["Content-Type"]

		cf, err := s.newCodec(ct)
		// TODO: needs better error handling
		if err != nil {
			sock.Send(&transport.Message{
				Header: map[string]string{
					"Content-Type": "text/plain",
				},
				Body: []byte(err.Error()),
			})
			s.wg.Done()
			return
		}

		codec := newRpcPlusCodec(&msg, sock, cf)

		// strip our headers
		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		delete(hdr, "Timeout")

		// set local/remote ips
		hdr["Local"] = sock.Local()
		hdr["Remote"] = sock.Remote()

		ctx := metadata.NewContext(context.Background(), hdr)

		// set the timeout if we have it
		if len(to) > 0 {
			if n, err := strconv.ParseUint(to, 10, 64); err == nil {
				ctx, _ = context.WithTimeout(ctx, time.Duration(n))
			}
		}

		// TODO: needs better error handling
		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
			s.wg.Done()
			log.Logf("Unexpected error serving request, closing socket: %v", err)
			return
		}
		s.wg.Done()
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	// update internal server
	s.rpc = &server{
		name:         s.opts.Name,
		serviceMap:   s.rpc.serviceMap,
		hdlrWrappers: s.opts.HdlrWrappers,
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}, opts ...HandlerOption) Handler {
	return newRpcHandler(h, opts...)
}

func (s *rpcServer) Handle(h Handler) error {
	s.Lock()
	defer s.Unlock()

	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}

	s.handlers[h.Name()] = h

	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}, opts ...SubscriberOption) Subscriber {
	return newSubscriber(topic, sb, opts...)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()
	node.Metadata["registry"] = config.Registry.String()

	s.RLock()
	// Maps are ordered randomly, sort the keys for consistency
	var handlerList []string
	for n, e := range s.handlers {
		// Only advertise non internal handlers
		if !e.Options().Internal {
			handlerList = append(handlerList, n)
		}
	}
	sort.Strings(handlerList)

	var subscriberList []*subscriber
	for e := range s.subscribers {
		// Only advertise non internal subscribers
		if !e.Options().Internal {
			subscriberList = append(subscriberList, e)
		}
	}
	sort.Slice(subscriberList, func(i, j int) bool {
		return subscriberList[i].topic > subscriberList[j].topic
	})

	var endpoints []*registry.Endpoint
	for _, n := range handlerList {
		endpoints = append(endpoints, s.handlers[n].Endpoints()...)
	}
	for _, e := range subscriberList {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	s.Lock()
	registered := s.registered
	s.Unlock()

	if !registered {
		log.Logf("Registering node: %s", node.Id)
	}

	// create registry options
	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}

	if err := config.Registry.Register(service, rOpts...); err != nil {
		return err
	}

	// already registered? don't need to register subscribers
	if registered {
		return nil
	}

	s.Lock()
	defer s.Unlock()

	s.registered = true

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		var opts []broker.SubscribeOption
		if queue := sb.Options().Queue; len(queue) > 0 {
			opts = append(opts, broker.Queue(queue))
		}
		sub, err := config.Broker.Subscribe(sb.Topic(), handler, opts...)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Logf("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()

	if !s.registered {
		s.Unlock()
		return nil
	}

	s.registered = false

	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Logf("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}

	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerDebugHandler(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Logf("Listening on %s", ts.Addr())
	s.Lock()
	// swap address
	addr := s.opts.Address
	s.opts.Address = ts.Addr()
	s.Unlock()

	exit := make(chan bool, 1)

	go func() {
		for {
			err := ts.Accept(s.accept)

			// check if we're supposed to exit
			select {
			case <-exit:
				return
			default:
			}

			// check the error and backoff
			if err != nil {
				log.Logf("Accept error: %v", err)
				time.Sleep(time.Second)
				continue
			}

			// no error just exit
			return
		}
	}()

	go func() {
		// wait for exit
		ch := <-s.exit
		exit <- true

		// wait for requests to finish
		if wait(s.opts.Context) {
			s.wg.Wait()
		}

		// close transport listener
		ch <- ts.Close()

		// disconnect the broker
		config.Broker.Disconnect()

		s.Lock()
		// swap back address
		s.opts.Address = addr
		s.Unlock()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	"github.com/trackit/jsonlog"
	"github.com/trackit/trackit/es"

	"github.com/trackit/trackit/awsSession"
	"github.com/trackit/trackit/config"
)

type (
	MessageData struct {
		TaskName   string   `json:"task_name"`
		Parameters []string `json:"parameters"`
		LogStream  string   `json:"cloudwatch_log_stream"`
	}

	LogInput struct {
		Message   string
		Timestamp int64
	}
)

const visibilityTimeoutInMinutes = 15
const waitTimeInSeconds = 20
const esHealthcheckTimeoutInMinutes = 5

func taskWorker(ctx context.Context) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger.Info("Running task 'worker'.", nil)
	sqsq := sqs.New(awsSession.Session)
	cwl := cloudwatchlogs.New(awsSession.Session)

	logger.Info("Retrieving SQS Queue URL.", nil)
	queueUrl, err := retrieveQueueUrl(ctx, sqsq)

	if err != nil {
		return err
	}

	logger.Info("Waiting for messages.", nil)

	var logsBuffer bytes.Buffer
	writer := io.MultiWriter(os.Stdout, &logsBuffer)
	ctx = jsonlog.ContextWithLogger(ctx, logger.WithWriter(writer))
	logger = jsonlog.LoggerFromContextOrDefault(ctx)

	for {
		message, receiptHandle, err := getNextMessage(ctx, sqsq, queueUrl)
		if err != nil {
			logsBuffer.Reset()
			continue
		}

		logger.Info("Received message, checking ES health and executing task.", map[string]interface{}{
			"message": message,
		})

		if res, err := es.Client.ClusterHealth().Timeout(strconv.Itoa(esHealthcheckTimeoutInMinutes) + "m").Do(ctx); err != nil || res.TimedOut {
			logger.Error("ES is not reachable.", map[string]interface{}{
				"timeout": strconv.Itoa(esHealthcheckTimeoutInMinutes) + "m",
				"error":   err.Error(),
			})
			logsBuffer.Reset()
			_ = resetMessageTimeout(ctx, sqsq, queueUrl, receiptHandle)
			continue
		}

		ctx = context.WithValue(ctx, "taskParameters", message.Parameters)

		if task, ok := tasks[message.TaskName]; ok {
			err = task(ctx)
			if err != nil {
				logger.Error("Error while executing task. Resetting visibility timeout to 0.", map[string]interface{}{
					"message": message,
					"error":   err.Error(),
				})
				_ = resetMessageTimeout(ctx, sqsq, queueUrl, receiptHandle)
				_ = flushCloudwatchLogEvents(ctx, cwl, message, &logsBuffer)
				continue
			}
			logger.Info("Task done, acknowledging.", nil)
			_ = acknowledgeMessage(ctx, sqsq, queueUrl, receiptHandle)
			_ = flushCloudwatchLogEvents(ctx, cwl, message, &logsBuffer)
		} else {
			logger.Error("Unable to find requested task.", map[string]interface{}{
				"task_name": message.TaskName,
			})
			logsBuffer.Reset()
			_ = acknowledgeMessage(ctx, sqsq, queueUrl, receiptHandle)
		}
	}
}

func getNextMessage(ctx context.Context, svc *sqs.SQS, queueUrl *string) (MessageData, *string, error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	messageData := MessageData{}
	msgResult, err := svc.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            queueUrl,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(60 * visibilityTimeoutInMinutes),
		WaitTimeSeconds:     aws.Int64(waitTimeInSeconds),
	})

	if err != nil {
		logger.Error("An error occurred while waiting for message.", map[string]interface{}{
			"queueUrl": *queueUrl,
			"error":    err.Error(),
		})
		return messageData, nil, err
	}
	if len(msgResult.Messages) == 0 {
		return messageData, nil, errors.New("queue: no message in queue.")
	}
	err = json.Unmarshal([]byte(*msgResult.Messages[0].Body), &messageData)
	if err != nil {
		logger.Error("Unable to decode message.", map[string]interface{}{
			"messageBody": *msgResult.Messages[0].Body,
			"error":       err.Error(),
		})
		return messageData, nil, err
	}
	return messageData, msgResult.Messages[0].ReceiptHandle, nil
}

func flushCloudwatchLogEvents(ctx context.Context, cwl *cloudwatchlogs.CloudWatchLogs, message MessageData, logsBuffer *bytes.Buffer) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)

	if config.Environment != "prod" && config.Environment != "stg" {
		logsBuffer.Reset()
		return nil
	}

	logGroup := config.Environment + "/task-logs/" + message.TaskName
	logStream := message.LogStream + "/" + uuid.New().String()

	_, err := cwl.CreateLogStream(&cloudwatchlogs.CreateLogStreamInput{
		LogGroupName:  aws.String(logGroup),
		LogStreamName: aws.String(logStream),
	})
	if err != nil {
		logger.Error("Unable to create log stream. Skipping logs sending.", map[string]interface{}{
			"logGroupName":  logGroup,
			"logStreamName": logStream,
			"error":         err.Error(),
		})
		logsBuffer.Reset()
		return err
	}

	logInputs := decodeLogBuffer(*logsBuffer)

	var logEvents []*cloudwatchlogs.InputLogEvent
	for _, logInput := range logInputs {
		fmt.Println(logInput.Message)
		logEvents = append(logEvents, &cloudwatchlogs.InputLogEvent{
			Message:   aws.String(logInput.Message),
			Timestamp: aws.Int64(logInput.Timestamp),
		})
	}

	_, err = cwl.PutLogEvents(&cloudwatchlogs.PutLogEventsInput{
		LogGroupName:  aws.String(logGroup),
		LogStreamName: aws.String(logStream),
		LogEvents:     logEvents,
	})

	if err != nil {
		logger.Error("Unable to put log stream events.", map[string]interface{}{
			"logGroupName":  logGroup,
			"logStreamName": logStream,
			"error":         err.Error(),
		})
		logsBuffer.Reset()
		return err
	}

	logsBuffer.Reset()
	return nil
}

func decodeLogBuffer(logsBuffer bytes.Buffer) []LogInput {
	lines := strings.Split(logsBuffer.String(), "\n")

	var logInputs []LogInput

	for _, line := range lines {
		if len(line) == 0 {
			continue
		}

		logInputs = append(logInputs, LogInput{
			Message:   line,
			Timestamp: getLogTimestamp(line),
		})
	}

	return logInputs
}

func getLogTimestamp(message string) int64 {
	var logJson map[string]interface{}

	err := json.Unmarshal([]byte(message), &logJson)
	if err != nil {
		return time.Now().UnixNano() / int64(time.Millisecond)
	}

	logTime, err := time.Parse(time.RFC3339, logJson["time"].(string))
	if err != nil {
		return time.Now().UnixNano() / int64(time.Millisecond)
	}

	return logTime.UnixNano() / int64(time.Millisecond)
}

func resetMessageTimeout(ctx context.Context, svc *sqs.SQS, queueUrl *string, receiptHandle *string) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	_, err := svc.ChangeMessageVisibility(&sqs.ChangeMessageVisibilityInput{
		ReceiptHandle:     receiptHandle,
		QueueUrl:          queueUrl,
		VisibilityTimeout: aws.Int64(0),
	})

	svc.SendMessage(&sqs.SendMessageInput{})
	if err != nil {
		logger.Error("Unable to reset timeout.", map[string]interface{}{
			"messageHandle": *receiptHandle,
			"error":         err.Error(),
		})
		return err
	}
	return nil
}

func acknowledgeMessage(ctx context.Context, svc *sqs.SQS, queueUrl *string, receiptHandle *string) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	_, err := svc.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      queueUrl,
		ReceiptHandle: receiptHandle,
	})
	if err != nil {
		logger.Error("Unable to acknowledge message.", map[string]interface{}{
			"messageHandle": *receiptHandle,
			"error":         err.Error(),
		})
		return err
	}

	return nil
}

func retrieveQueueUrl(ctx context.Context, svc *sqs.SQS) (queueUrl *string, err error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	urlResult, err := svc.GetQueueUrl(&sqs.GetQueueUrlInput{
		QueueName: &config.SQSQueueName,
	})
	if err != nil {
		logger.Error("Unable to get queue URL from name.", map[string]interface{}{
			"queueName": config.SQSQueueName,
			"error":     err.Error(),
		})
		return nil, err
	}
	return urlResult.QueueUrl, nil
}

func paramsFromContextOrArgs(ctx context.Context) []string {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	args, ok := ctx.Value("taskParameters").([]string)
	if !ok {
		logger.Info("Retrieving task parameters from args as context does not contain any parameters.", nil)
		args = flag.Args()
	}
	return args
}
<commit_msg>Changed visibility timeout and ES check timeout<commit_after>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	"github.com/trackit/jsonlog"
	"github.com/trackit/trackit/es"

	"github.com/trackit/trackit/awsSession"
	"github.com/trackit/trackit/config"
)

type (
	MessageData struct {
		TaskName   string   `json:"task_name"`
		Parameters []string `json:"parameters"`
		LogStream  string   `json:"cloudwatch_log_stream"`
	}

	LogInput struct {
		Message   string
		Timestamp int64
	}
)

const visibilityTimeoutInMinutes = 30
const waitTimeInSeconds = 20
const esHealthcheckTimeoutInMinutes = 10

func taskWorker(ctx context.Context) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger.Info("Running task 'worker'.", nil)
	sqsq := sqs.New(awsSession.Session)
	cwl := cloudwatchlogs.New(awsSession.Session)

	logger.Info("Retrieving SQS Queue URL.", nil)
	queueUrl, err := retrieveQueueUrl(ctx, sqsq)

	if err != nil {
		return err
	}

	logger.Info("Waiting for messages.", nil)

	var logsBuffer bytes.Buffer
	writer := io.MultiWriter(os.Stdout, &logsBuffer)
	ctx = jsonlog.ContextWithLogger(ctx, logger.WithWriter(writer))
	logger = jsonlog.LoggerFromContextOrDefault(ctx)

	for {
		message, receiptHandle, err := getNextMessage(ctx, sqsq, queueUrl)
		if err != nil {
			logsBuffer.Reset()
			continue
		}

		logger.Info("Received message, checking ES health and executing task.", map[string]interface{}{
			"message": message,
		})

		if res, err := es.Client.ClusterHealth().Timeout(strconv.Itoa(esHealthcheckTimeoutInMinutes) + "m").Do(ctx); err != nil || res.TimedOut {
			logger.Error("ES is not reachable.", map[string]interface{}{
				"timeout": strconv.Itoa(esHealthcheckTimeoutInMinutes) + "m",
				"error":   err.Error(),
			})
			logsBuffer.Reset()
			_ = resetMessageTimeout(ctx, sqsq, queueUrl, receiptHandle)
			continue
		}

		ctx = context.WithValue(ctx, "taskParameters", message.Parameters)

		if task, ok := tasks[message.TaskName]; ok {
			err = task(ctx)
			if err != nil {
				logger.Error("Error while executing task. Resetting visibility timeout to 0.", map[string]interface{}{
					"message": message,
					"error":   err.Error(),
				})
				_ = resetMessageTimeout(ctx, sqsq, queueUrl, receiptHandle)
				_ = flushCloudwatchLogEvents(ctx, cwl, message, &logsBuffer)
				continue
			}
			logger.Info("Task done, acknowledging.", nil)
			_ = acknowledgeMessage(ctx, sqsq, queueUrl, receiptHandle)
			_ = flushCloudwatchLogEvents(ctx, cwl, message, &logsBuffer)
		} else {
			logger.Error("Unable to find requested task.", map[string]interface{}{
				"task_name": message.TaskName,
			})
			logsBuffer.Reset()
			_ = acknowledgeMessage(ctx, sqsq, queueUrl, receiptHandle)
		}
	}
}

func getNextMessage(ctx context.Context, svc *sqs.SQS, queueUrl *string) (MessageData, *string, error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	messageData := MessageData{}
	msgResult, err := svc.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            queueUrl,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(60 * visibilityTimeoutInMinutes),
		WaitTimeSeconds:     aws.Int64(waitTimeInSeconds),
	})

	if err != nil {
		logger.Error("An error occurred while waiting for message.", map[string]interface{}{
			"queueUrl": *queueUrl,
			"error":    err.Error(),
		})
		return messageData, nil, err
	}
	if len(msgResult.Messages) == 0 {
		return messageData, nil, errors.New("queue: no message in queue.")
	}
	err = json.Unmarshal([]byte(*msgResult.Messages[0].Body), &messageData)
	if err != nil {
		logger.Error("Unable to decode message.", map[string]interface{}{
			"messageBody": *msgResult.Messages[0].Body,
			"error":       err.Error(),
		})
		return messageData, nil, err
	}
	return messageData, msgResult.Messages[0].ReceiptHandle, nil
}

func flushCloudwatchLogEvents(ctx context.Context, cwl *cloudwatchlogs.CloudWatchLogs, message MessageData, logsBuffer *bytes.Buffer) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)

	if config.Environment != "prod" && config.Environment != "stg" {
		logsBuffer.Reset()
		return nil
	}

	logGroup := config.Environment + "/task-logs/" + message.TaskName
	logStream := message.LogStream + "/" + uuid.New().String()

	_, err := cwl.CreateLogStream(&cloudwatchlogs.CreateLogStreamInput{
		LogGroupName:  aws.String(logGroup),
		LogStreamName: aws.String(logStream),
	})
	if err != nil {
		logger.Error("Unable to create log stream. Skipping logs sending.", map[string]interface{}{
			"logGroupName":  logGroup,
			"logStreamName": logStream,
			"error":         err.Error(),
		})
		logsBuffer.Reset()
		return err
	}

	logInputs := decodeLogBuffer(*logsBuffer)

	var logEvents []*cloudwatchlogs.InputLogEvent
	for _, logInput := range logInputs {
		fmt.Println(logInput.Message)
		logEvents = append(logEvents, &cloudwatchlogs.InputLogEvent{
			Message:   aws.String(logInput.Message),
			Timestamp: aws.Int64(logInput.Timestamp),
		})
	}

	_, err = cwl.PutLogEvents(&cloudwatchlogs.PutLogEventsInput{
		LogGroupName:  aws.String(logGroup),
		LogStreamName: aws.String(logStream),
		LogEvents:     logEvents,
	})

	if err != nil {
		logger.Error("Unable to put log stream events.", map[string]interface{}{
			"logGroupName":  logGroup,
			"logStreamName": logStream,
			"error":         err.Error(),
		})
		logsBuffer.Reset()
		return err
	}

	logsBuffer.Reset()
	return nil
}

func decodeLogBuffer(logsBuffer bytes.Buffer) []LogInput {
	lines := strings.Split(logsBuffer.String(), "\n")

	var logInputs []LogInput

	for _, line := range lines {
		if len(line) == 0 {
			continue
		}

		logInputs = append(logInputs, LogInput{
			Message:   line,
			Timestamp: getLogTimestamp(line),
		})
	}

	return logInputs
}

func getLogTimestamp(message string) int64 {
	var logJson map[string]interface{}

	err := json.Unmarshal([]byte(message), &logJson)
	if err != nil {
		return time.Now().UnixNano() / int64(time.Millisecond)
	}

	logTime, err := time.Parse(time.RFC3339, logJson["time"].(string))
	if err != nil {
		return time.Now().UnixNano() / int64(time.Millisecond)
	}

	return logTime.UnixNano() / int64(time.Millisecond)
}

func resetMessageTimeout(ctx context.Context, svc *sqs.SQS, queueUrl *string, receiptHandle *string) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	_, err := svc.ChangeMessageVisibility(&sqs.ChangeMessageVisibilityInput{
		ReceiptHandle:     receiptHandle,
		QueueUrl:          queueUrl,
		VisibilityTimeout: aws.Int64(0),
	})

	svc.SendMessage(&sqs.SendMessageInput{})
	if err != nil {
		logger.Error("Unable to reset timeout.", map[string]interface{}{
			"messageHandle": *receiptHandle,
			"error":         err.Error(),
		})
		return err
	}
	return nil
}

func acknowledgeMessage(ctx context.Context, svc *sqs.SQS, queueUrl *string, receiptHandle *string) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	_, err := svc.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      queueUrl,
		ReceiptHandle: receiptHandle,
	})
	if err != nil {
		logger.Error("Unable to acknowledge message.", map[string]interface{}{
			"messageHandle": *receiptHandle,
			"error":         err.Error(),
		})
		return err
	}

	return nil
}

func retrieveQueueUrl(ctx context.Context, svc *sqs.SQS) (queueUrl *string, err error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	urlResult, err := svc.GetQueueUrl(&sqs.GetQueueUrlInput{
		QueueName: &config.SQSQueueName,
	})
	if err != nil {
		logger.Error("Unable to get queue URL from name.", map[string]interface{}{
			"queueName": config.SQSQueueName,
			"error":     err.Error(),
		})
		return nil, err
	}
	return urlResult.QueueUrl, nil
}

func paramsFromContextOrArgs(ctx context.Context) []string {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	args, ok := ctx.Value("taskParameters").([]string)
	if !ok {
		logger.Info("Retrieving task parameters from args as context does not contain any parameters.", nil)
		args = flag.Args()
	}
	return args
}
<|endoftext|>


<commit_before>package rrserver

import (
	"github.com/songtianyi/rrframework/logs"
	"io"
	"net"
	"strconv"
)

var (
	CustomHandleConn = func(c *TCPConnection, packet []byte) {
		logs.Warn("forget rrserver.CustomHandleConn = YourCustomHandleConn in init func?")
	}
)

type TCPServer struct {
	ls   net.Listener
	port int
}

func CreateTCPServer(inf string, port int) (error, *TCPServer) {
	err, ipaddr := getIpAddrByInterface(inf)
	if err != nil {
		return err, nil
	}
	listener, err := net.Listen("tcp", net.JoinHostPort(ipaddr, strconv.Itoa(port)))
	if err != nil {
		return err, nil
	}
	s := &TCPServer{
		ls:   listener,
		port: port,
	}
	return nil, s
}

func CreateTCPServerByAddr(addr string) (error, *TCPServer) {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err, nil
	}
	s := &TCPServer{
		ls:   listener,
		port: port,
	}
	return nil, s
}

func (s *TCPServer) Start() {
	logs.Info("Server listening in [%s]", s.ls.Addr())
	for {
		conn, err := s.ls.Accept()
		if err != nil {
			logs.Error("Server Accept() return error, %s", err)
			break
		}
		go s.handleConn(NewTCPConnection(conn))
	}
	return
}

func (s *TCPServer) handleConn(c *TCPConnection) {
	for {
		err, packet := c.Read()
		if err != nil {
			// end goroutine
			if err != io.EOF {
				logs.Error("Error occurred when read packet, %s", err)
			}
			return
		}
		// Maybe thousands of packet coming in the same time
		// So lock for connection is necessary
		go CustomHandleConn(c, packet)
	}
}
<commit_msg>fix: not set port<commit_after>package rrserver

import (
	"github.com/songtianyi/rrframework/logs"
	"io"
	"net"
	"strconv"
)

var (
	CustomHandleConn = func(c *TCPConnection, packet []byte) {
		logs.Warn("forget rrserver.CustomHandleConn = YourCustomHandleConn in init func?")
	}
)

type TCPServer struct {
	ls   net.Listener
	port int
}

func CreateTCPServer(inf string, port int) (error, *TCPServer) {
	err, ipaddr := getIpAddrByInterface(inf)
	if err != nil {
		return err, nil
	}
	listener, err := net.Listen("tcp", net.JoinHostPort(ipaddr, strconv.Itoa(port)))
	if err != nil {
		return err, nil
	}
	s := &TCPServer{
		ls:   listener,
		port: port,
	}
	return nil, s
}

func CreateTCPServerByAddr(addr string) (error, *TCPServer) {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err, nil
	}
	parts := strings.Split(addr, ":")
	if len(parts) != 2 {
		return fmt.Errorf("addr invalid len=%d", len(parts))
	}
	p, err := strconv.Atoi(parts[1])
	if err != nil {
		return err
	}
	s := &TCPServer{
		ls:   listener,
		port: p,
	}
	return nil, s
}

func (s *TCPServer) Start() {
	logs.Info("Server listening in [%s]", s.ls.Addr())
	for {
		conn, err := s.ls.Accept()
		if err != nil {
			logs.Error("Server Accept() return error, %s", err)
			break
		}
		go s.handleConn(NewTCPConnection(conn))
	}
	return
}

func (s *TCPServer) handleConn(c *TCPConnection) {
	for {
		err, packet := c.Read()
		if err != nil {
			// end goroutine
			if err != io.EOF {
				logs.Error("Error occurred when read packet, %s", err)
			}
			return
		}
		// Maybe thousands of packet coming in the same time
		// So lock for connection is necessary
		go CustomHandleConn(c, packet)
	}
}
<|endoftext|>


<commit_before>package inputviz

import (
	"fmt"
	"image/color"

	"github.com/oakmound/oak/v3/alg/floatgeom"
	"github.com/oakmound/oak/v3/event"
	"github.com/oakmound/oak/v3/key"
	"github.com/oakmound/oak/v3/render"
	"github.com/oakmound/oak/v3/scene"
)

type KeyboardLayout interface {
	KeyRect(k string) floatgeom.Rect2
}

type LayoutKey interface {
	Pos() LayoutPosition
}

type LayoutPosition struct {
	Key    string
	Gap    bool
	Width  float64
	Height float64
	Row    float64
	Col    float64
}

type gap float64

func (g gap) Pos() LayoutPosition {
	return LayoutPosition{
		Gap:   true,
		Width: float64(g),
	}
}

type standardKey string

func (s standardKey) Pos() LayoutPosition {
	return LayoutPosition{
		Key:    string(s),
		Width:  1,
		Height: 1,
	}
}

type wideKey struct {
	k string
	w float64
}

func (w wideKey) Pos() LayoutPosition {
	return LayoutPosition{
		Key:    string(w.k),
		Width:  w.w,
		Height: 1,
	}
}

type tallKey struct {
	k string
	h float64
}

func (h tallKey) Pos() LayoutPosition {
	return LayoutPosition{
		Key:    string(h.k),
		Width:  1,
		Height: h.h,
	}
}

type LayoutQWERTY struct {
	Bounds    floatgeom.Rect2
	layoutMap map[string]LayoutPosition
}

func (l *LayoutQWERTY) init() {
	if l.layoutMap != nil {
		return
	}
	type sk = standardKey

	l.layoutMap = make(map[string]LayoutPosition)
	qwertyRows := [][]LayoutKey{
		{sk(key.Escape), gap(1), sk(key.F1), sk(key.F2), sk(key.F3), sk(key.F4), gap(.5), sk(key.F5), sk(key.F6), sk(key.F7), sk(key.F8), gap(.5), sk(key.F9), sk(key.F10), sk(key.F11), sk(key.F12), gap(2.1), sk(key.Pause)},
		{gap(1), sk(key.One), sk(key.Two), sk(key.Three), sk(key.Four), sk(key.Five), sk(key.Six), sk(key.Seven), sk(key.Eight), sk(key.Nine), sk(key.Zero), sk(key.HyphenMinus), sk(key.EqualSign), wideKey{key.DeleteBackspace, 2.0}, gap(.1), sk(key.Insert), sk(key.Home), sk(key.PageUp), gap(.1), sk(key.KeypadNumLock), sk(key.KeypadSlash), sk(key.KeypadAsterisk), sk(key.KeypadHyphenMinus)},
		{wideKey{key.Tab, 1.5}, sk(key.Q), sk(key.W), sk(key.E), sk(key.R), sk(key.T), sk(key.Y), sk(key.U), sk(key.I), sk(key.O), sk(key.P), sk(key.LeftSquareBracket), sk(key.RightSquareBracket), wideKey{key.Backslash, 1.5}, gap(.1), sk(key.DeleteForward), sk(key.End), sk(key.PageDown), gap(.1), sk(key.Keypad7), sk(key.Keypad8), sk(key.Keypad9), tallKey{key.KeypadPlusSign, 2}},
		{wideKey{key.CapsLock, 1.5}, sk(key.A), sk(key.S), sk(key.D), sk(key.F), sk(key.G), sk(key.H), sk(key.J), sk(key.K), sk(key.L), sk(key.Semicolon), sk(key.Apostrophe), wideKey{key.ReturnEnter, 2.5}, gap(3.2), sk(key.Keypad4), sk(key.Keypad5), sk(key.Keypad6)},
		{wideKey{key.LeftShift, 2.0}, sk(key.Z), sk(key.X), sk(key.C), sk(key.V), sk(key.B), sk(key.N), sk(key.M), sk(key.Comma), sk(key.FullStop), sk(key.Slash), wideKey{key.RightShift, 3.0}, gap(1.1), sk(key.UpArrow), gap(1.1), sk(key.Keypad1), sk(key.Keypad2), sk(key.Keypad3), tallKey{key.KeypadEnter, 2.0}},
		{wideKey{key.LeftControl, 1.5}, sk(key.LeftGUI), wideKey{key.LeftAlt, 1.5}, wideKey{key.Spacebar, 7.0}, wideKey{key.RightAlt, 1.5}, sk(key.RightGUI), wideKey{key.RightControl, 1.5}, gap(.1), sk(key.LeftArrow), sk(key.DownArrow), sk(key.RightArrow), gap(.1), wideKey{key.Keypad0, 2.0}, sk(key.KeypadPeriod)},
	}
	rowFloats := []float64{0.0, 1.1, 2.1, 3.1, 4.1, 5.1}
	for row, cols := range qwertyRows {
		rf := rowFloats[row]
		cf := 0.0
		for _, v := range cols {
			ps := v.Pos()
			if ps.Key != "" {
				l.layoutMap[ps.Key] = LayoutPosition{
					Row:    rf,
					Col:    cf,
					Width:  ps.Width - .1,
					Height: ps.Height - .1,
				}
			}
			cf += ps.Width
		}
	}
}

func (l *LayoutQWERTY) KeyRect(k string) floatgeom.Rect2 {
	l.init()

	pos, ok := l.layoutMap[k]
	if !ok {
		return floatgeom.Rect2{}
	}
	row := pos.Row
	col := pos.Col
	width := pos.Width
	height := pos.Height

	w, h := l.Bounds.W(), l.Bounds.H()
	// max row = 5.1, drawn down to 6
	// max col = 21.2, drawn right to 22.1
	rowHeight := h / 6.0
	colWidth := w / 22.1

	x := col * colWidth
	y := row * rowHeight
	keyHeight := height * rowHeight
	keyWidth := width * colWidth

	return floatgeom.NewRect2WH(x, y, keyWidth, keyHeight)
}

var defaultColors = map[string]color.Color{}

type Keyboard struct {
	Rect      floatgeom.Rect2
	BaseLayer int
	Colors    map[string]color.Color
	KeyboardLayout

	RenderCharacters bool
	Font             *render.Font

	event.CID
	ctx *scene.Context

	rs map[string]*render.Switch
}

func (k *Keyboard) Init() event.CID {
	k.CID = k.ctx.CallerMap.NextID(k)
	return k.CID
}

func (k *Keyboard) RenderAndListen(ctx *scene.Context, layer int) error {
	k.ctx = ctx
	k.Init()

	if k.Rect.W() == 0 || k.Rect.H() == 0 {
		k.Rect.Max = k.Rect.Min.Add(floatgeom.Point2{320, 180})
	}
	if k.KeyboardLayout == nil {
		k.KeyboardLayout = &LayoutQWERTY{
			Bounds: k.Rect,
		}
	}
	if k.Colors == nil {
		k.Colors = defaultColors
	}
	if k.Font == nil {
		k.Font = render.DefaultFont()
	}

	k.rs = make(map[string]*render.Switch)

	for kv := range key.AllKeys {
		rect := k.KeyboardLayout.KeyRect(kv)
		if rect == (floatgeom.Rect2{}) {
			continue
		}
		pressedColor := color.RGBA{255, 255, 255, 255}
		var unpressedColor color.Color = color.RGBA{160, 160, 160, 255}
		if c, ok := k.Colors[kv]; ok {
			unpressedColor = c
		}
		r := render.NewSwitch("released", map[string]render.Modifiable{
			"pressed":  render.NewColorBox(int(rect.W()), int(rect.H()), pressedColor),
			"released": render.NewColorBox(int(rect.W()), int(rect.H()), unpressedColor),
		})
		r.SetPos(rect.Min.X(), rect.Min.Y())
		k.rs[kv] = r
		if k.RenderCharacters {
			fmt.Println("kv:", kv)
			x, y := rect.Min.X(), rect.Min.Y()
			txt := k.Font.NewText(kv, x, y)
			tw, th := txt.GetDims()
			xBuffer := rect.W() - float64(tw)
			yBuffer := rect.H() - float64(th)
			// Only render strings that will stay inside their boundaries
			if xBuffer >= 0 {
				txt.ShiftX(xBuffer / 2)
				txt.ShiftY(yBuffer / 2)
				if k.BaseLayer == -1 {
					ctx.DrawStack.Draw(txt, layer+1)
				} else {
					ctx.DrawStack.Draw(txt, k.BaseLayer, layer+1)
				}
			}
		}
		if k.BaseLayer == -1 {
			ctx.DrawStack.Draw(r, layer)
		} else {
			ctx.DrawStack.Draw(r, k.BaseLayer, layer)
		}
	}

	k.Bind(key.Down, key.Binding(func(id event.CID, ev key.Event) int {
		kb, _ := k.ctx.CallerMap.GetEntity(id).(*Keyboard)
		btn := ev.Code.String()[4:]
		if kb.rs[btn] == nil {
			return 0
		}
		kb.rs[btn].Set("pressed")
		return 0
	}))
	k.Bind(key.Up, key.Binding(func(id event.CID, ev key.Event) int {
		kb, _ := k.ctx.CallerMap.GetEntity(id).(*Keyboard)
		btn := ev.Code.String()[4:]
		if kb.rs[btn] == nil {
			return 0
		}
		kb.rs[btn].Set("released")
		return 0
	}))

	return nil
}

func (k *Keyboard) Destroy() {
	k.UnbindAll()
	for _, r := range k.rs {
		r.Undraw()
	}
}
<commit_msg>debugtools/inputviz: include grave accent key<commit_after>package inputviz

import (
	"image/color"

	"github.com/oakmound/oak/v3/alg/floatgeom"
	"github.com/oakmound/oak/v3/event"
	"github.com/oakmound/oak/v3/key"
	"github.com/oakmound/oak/v3/render"
	"github.com/oakmound/oak/v3/scene"
)

type KeyboardLayout interface {
	KeyRect(k string) floatgeom.Rect2
}

type LayoutKey interface {
	Pos() LayoutPosition
}

type LayoutPosition struct {
	Key    string
	Gap    bool
	Width  float64
	Height float64
	Row    float64
	Col    float64
}

type gap float64

func (g gap) Pos() LayoutPosition {
	return LayoutPosition{
		Gap:   true,
		Width: float64(g),
	}
}

type standardKey string

func (s standardKey) Pos() LayoutPosition {
	return LayoutPosition{
		Key:    string(s),
		Width:  1,
		Height: 1,
	}
}

type wideKey struct {
	k string
	w float64
}

func (w wideKey) Pos() LayoutPosition {
	return LayoutPosition{
		Key:    string(w.k),
		Width:  w.w,
		Height: 1,
	}
}

type tallKey struct {
	k string
	h float64
}

func (h tallKey) Pos() LayoutPosition {
	return LayoutPosition{
		Key:    string(h.k),
		Width:  1,
		Height: h.h,
	}
}

type LayoutQWERTY struct {
	Bounds    floatgeom.Rect2
	layoutMap map[string]LayoutPosition
}

func (l *LayoutQWERTY) init() {
	if l.layoutMap != nil {
		return
	}
	type sk = standardKey

	l.layoutMap = make(map[string]LayoutPosition)
	qwertyRows := [][]LayoutKey{
		{sk(key.Escape), gap(1), sk(key.F1), sk(key.F2), sk(key.F3), sk(key.F4), gap(.5), sk(key.F5), sk(key.F6), sk(key.F7), sk(key.F8), gap(.5), sk(key.F9), sk(key.F10), sk(key.F11), sk(key.F12), gap(2.1), sk(key.Pause)},
		{sk(key.GraveAccent), sk(key.One), sk(key.Two), sk(key.Three), sk(key.Four), sk(key.Five), sk(key.Six), sk(key.Seven), sk(key.Eight), sk(key.Nine), sk(key.Zero), sk(key.HyphenMinus), sk(key.EqualSign), wideKey{key.DeleteBackspace, 2.0}, gap(.1), sk(key.Insert), sk(key.Home), sk(key.PageUp), gap(.1), sk(key.KeypadNumLock), sk(key.KeypadSlash), sk(key.KeypadAsterisk), sk(key.KeypadHyphenMinus)},
		{wideKey{key.Tab, 1.5}, sk(key.Q), sk(key.W), sk(key.E), sk(key.R), sk(key.T), sk(key.Y), sk(key.U), sk(key.I), sk(key.O), sk(key.P), sk(key.LeftSquareBracket), sk(key.RightSquareBracket), wideKey{key.Backslash, 1.5}, gap(.1), sk(key.DeleteForward), sk(key.End), sk(key.PageDown), gap(.1), sk(key.Keypad7), sk(key.Keypad8), sk(key.Keypad9), tallKey{key.KeypadPlusSign, 2}},
		{wideKey{key.CapsLock, 1.5}, sk(key.A), sk(key.S), sk(key.D), sk(key.F), sk(key.G), sk(key.H), sk(key.J), sk(key.K), sk(key.L), sk(key.Semicolon), sk(key.Apostrophe), wideKey{key.ReturnEnter, 2.5}, gap(3.2), sk(key.Keypad4), sk(key.Keypad5), sk(key.Keypad6)},
		{wideKey{key.LeftShift, 2.0}, sk(key.Z), sk(key.X), sk(key.C), sk(key.V), sk(key.B), sk(key.N), sk(key.M), sk(key.Comma), sk(key.FullStop), sk(key.Slash), wideKey{key.RightShift, 3.0}, gap(1.1), sk(key.UpArrow), gap(1.1), sk(key.Keypad1), sk(key.Keypad2), sk(key.Keypad3), tallKey{key.KeypadEnter, 2.0}},
		{wideKey{key.LeftControl, 1.5}, sk(key.LeftGUI), wideKey{key.LeftAlt, 1.5}, wideKey{key.Spacebar, 7.0}, wideKey{key.RightAlt, 1.5}, sk(key.RightGUI), wideKey{key.RightControl, 1.5}, gap(.1), sk(key.LeftArrow), sk(key.DownArrow), sk(key.RightArrow), gap(.1), wideKey{key.Keypad0, 2.0}, sk(key.KeypadPeriod)},
	}
	rowFloats := []float64{0.0, 1.1, 2.1, 3.1, 4.1, 5.1}
	for row, cols := range qwertyRows {
		rf := rowFloats[row]
		cf := 0.0
		for _, v := range cols {
			ps := v.Pos()
			if ps.Key != "" {
				l.layoutMap[ps.Key] = LayoutPosition{
					Row:    rf,
					Col:    cf,
					Width:  ps.Width - .1,
					Height: ps.Height - .1,
				}
			}
			cf += ps.Width
		}
	}
}

func (l *LayoutQWERTY) KeyRect(k string) floatgeom.Rect2 {
	l.init()

	pos, ok := l.layoutMap[k]
	if !ok {
		return floatgeom.Rect2{}
	}
	row := pos.Row
	col := pos.Col
	width := pos.Width
	height := pos.Height

	w, h := l.Bounds.W(), l.Bounds.H()
	// max row = 5.1, drawn down to 6
	// max col = 21.2, drawn right to 22.1
	rowHeight := h / 6.0
	colWidth := w / 22.1

	x := col * colWidth
	y := row * rowHeight
	keyHeight := height * rowHeight
	keyWidth := width * colWidth

	return floatgeom.NewRect2WH(x, y, keyWidth, keyHeight)
}

var defaultColors = map[string]color.Color{}

type Keyboard struct {
	Rect      floatgeom.Rect2
	BaseLayer int
	Colors    map[string]color.Color
	KeyboardLayout

	RenderCharacters bool
	Font             *render.Font

	event.CID
	ctx *scene.Context

	rs map[string]*render.Switch
}

func (k *Keyboard) Init() event.CID {
	k.CID = k.ctx.CallerMap.NextID(k)
	return k.CID
}

func (k *Keyboard) RenderAndListen(ctx *scene.Context, layer int) error {
	k.ctx = ctx
	k.Init()

	if k.Rect.W() == 0 || k.Rect.H() == 0 {
		k.Rect.Max = k.Rect.Min.Add(floatgeom.Point2{320, 180})
	}
	if k.KeyboardLayout == nil {
		k.KeyboardLayout = &LayoutQWERTY{
			Bounds: k.Rect,
		}
	}
	if k.Colors == nil {
		k.Colors = defaultColors
	}
	if k.Font == nil {
		k.Font = render.DefaultFont()
	}

	k.rs = make(map[string]*render.Switch)

	for kv := range key.AllKeys {
		rect := k.KeyboardLayout.KeyRect(kv)
		if rect == (floatgeom.Rect2{}) {
			continue
		}
		pressedColor := color.RGBA{255, 255, 255, 255}
		var unpressedColor color.Color = color.RGBA{160, 160, 160, 255}
		if c, ok := k.Colors[kv]; ok {
			unpressedColor = c
		}
		r := render.NewSwitch("released", map[string]render.Modifiable{
			"pressed":  render.NewColorBox(int(rect.W()), int(rect.H()), pressedColor),
			"released": render.NewColorBox(int(rect.W()), int(rect.H()), unpressedColor),
		})
		r.SetPos(rect.Min.X(), rect.Min.Y())
		k.rs[kv] = r
		if k.RenderCharacters {
			x, y := rect.Min.X(), rect.Min.Y()
			txt := k.Font.NewText(kv, x, y)
			tw, th := txt.GetDims()
			xBuffer := rect.W() - float64(tw)
			yBuffer := rect.H() - float64(th)
			// Only render strings that will stay inside their boundaries
			if xBuffer >= 0 {
				txt.ShiftX(xBuffer / 2)
				txt.ShiftY(yBuffer / 2)
				if k.BaseLayer == -1 {
					ctx.DrawStack.Draw(txt, layer+1)
				} else {
					ctx.DrawStack.Draw(txt, k.BaseLayer, layer+1)
				}
			}
		}
		if k.BaseLayer == -1 {
			ctx.DrawStack.Draw(r, layer)
		} else {
			ctx.DrawStack.Draw(r, k.BaseLayer, layer)
		}
	}

	k.Bind(key.Down, key.Binding(func(id event.CID, ev key.Event) int {
		kb, _ := k.ctx.CallerMap.GetEntity(id).(*Keyboard)
		btn := ev.Code.String()[4:]
		if kb.rs[btn] == nil {
			return 0
		}
		kb.rs[btn].Set("pressed")
		return 0
	}))
	k.Bind(key.Up, key.Binding(func(id event.CID, ev key.Event) int {
		kb, _ := k.ctx.CallerMap.GetEntity(id).(*Keyboard)
		btn := ev.Code.String()[4:]
		if kb.rs[btn] == nil {
			return 0
		}
		kb.rs[btn].Set("released")
		return 0
	}))

	return nil
}

func (k *Keyboard) Destroy() {
	k.UnbindAll()
	for _, r := range k.rs {
		r.Undraw()
	}
}
<|endoftext|>


<commit_before>package web

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/registry"
	maddr "github.com/micro/go-micro/util/addr"
	mhttp "github.com/micro/go-micro/util/http"
	"github.com/micro/go-micro/util/log"
	mnet "github.com/micro/go-micro/util/net"
	mls "github.com/micro/go-micro/util/tls"
)

type service struct {
	opts Options

	mux *http.ServeMux
	srv *registry.Service

	sync.Mutex
	running bool
	static  bool
	exit    chan chan error
}

func newService(opts ...Option) Service {
	options := newOptions(opts...)
	s := &service{
		opts:   options,
		mux:    http.NewServeMux(),
		static: true,
	}
	s.srv = s.genSrv()
	return s
}

func (s *service) genSrv() *registry.Service {
	// default host:port
	parts := strings.Split(s.opts.Address, ":")
	host := strings.Join(parts[:len(parts)-1], ":")
	port, _ := strconv.Atoi(parts[len(parts)-1])

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(s.opts.Advertise) > 0 {
		parts = strings.Split(s.opts.Advertise, ":")

		// we have host:port
		if len(parts) > 1 {
			// set the host
			host = strings.Join(parts[:len(parts)-1], ":")

			// get the port
			if aport, _ := strconv.Atoi(parts[len(parts)-1]); aport > 0 {
				port = aport
			}
		} else {
			host = parts[0]
		}
	}

	addr, err := maddr.Extract(host)
	if err != nil {
		// best effort localhost
		addr = "127.0.0.1"
	}

	return &registry.Service{
		Name:    s.opts.Name,
		Version: s.opts.Version,
		Nodes: []*registry.Node{{
			Id:       s.opts.Id,
			Address:  fmt.Sprintf("%s:%d", addr, port),
			Metadata: s.opts.Metadata,
		}},
	}
}

func (s *service) run(exit chan bool) {
	if s.opts.RegisterInterval <= time.Duration(0) {
		return
	}

	t := time.NewTicker(s.opts.RegisterInterval)

	for {
		select {
		case <-t.C:
			s.register()
		case <-exit:
			t.Stop()
			return
		}
	}
}

func (s *service) register() error {
	if s.srv == nil {
		return nil
	}
	// default to service registry
	r := s.opts.Service.Client().Options().Registry
	// switch to option if specified
	if s.opts.Registry != nil {
		r = s.opts.Registry
	}

	// service node need modify, node address maybe changed
	srv := s.genSrv()
	srv.Endpoints = s.srv.Endpoints
	s.srv = srv
	return r.Register(s.srv, registry.RegisterTTL(s.opts.RegisterTTL))
}

func (s *service) deregister() error {
	if s.srv == nil {
		return nil
	}
	// default to service registry
	r := s.opts.Service.Client().Options().Registry
	// switch to option if specified
	if s.opts.Registry != nil {
		r = s.opts.Registry
	}
	return r.Deregister(s.srv)
}

func (s *service) start() error {
	s.Lock()
	defer s.Unlock()

	if s.running {
		return nil
	}

	l, err := s.listen("tcp", s.opts.Address)
	if err != nil {
		return err
	}

	s.opts.Address = l.Addr().String()
	srv := s.genSrv()
	srv.Endpoints = s.srv.Endpoints
	s.srv = srv

	var h http.Handler

	if s.opts.Handler != nil {
		h = s.opts.Handler
	} else {
		h = s.mux
		var r sync.Once

		// register the html dir
		r.Do(func() {
			// static dir
			static := s.opts.StaticDir
			if s.opts.StaticDir[0] != '/' {
				dir, _ := os.Getwd()
				static = filepath.Join(dir, static)
			}

			// set static if no / handler is registered
			if s.static {
				_, err := os.Stat(static)
				if err == nil {
					log.Logf("Enabling static file serving from %s", static)
					s.mux.Handle("/", http.FileServer(http.Dir(static)))
				}
			}
		})
	}

	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	var httpSrv *http.Server
	if s.opts.Server != nil {
		httpSrv = s.opts.Server
	} else {
		httpSrv = &http.Server{}
	}

	httpSrv.Handler = h

	go httpSrv.Serve(l)

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	s.exit = make(chan chan error, 1)
	s.running = true

	go func() {
		ch := <-s.exit
		ch <- l.Close()
	}()

	log.Logf("Listening on %v\n", l.Addr().String())
	return nil
}

func (s *service) stop() error {
	s.Lock()
	defer s.Unlock()

	if !s.running {
		return nil
	}

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			return err
		}
	}

	ch := make(chan error, 1)
	s.exit <- ch
	s.running = false

	log.Log("Stopping")

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			if chErr := <-ch; chErr != nil {
				return chErr
			}
			return err
		}
	}

	return <-ch
}

func (s *service) Client() *http.Client {
	rt := mhttp.NewRoundTripper(
		mhttp.WithRegistry(registry.DefaultRegistry),
	)
	return &http.Client{
		Transport: rt,
	}
}

func (s *service) Handle(pattern string, handler http.Handler) {
	var seen bool
	for _, ep := range s.srv.Endpoints {
		if ep.Name == pattern {
			seen = true
			break
		}
	}

	// if its unseen then add an endpoint
	if !seen {
		s.srv.Endpoints = append(s.srv.Endpoints, &registry.Endpoint{
			Name: pattern,
		})
	}

	// disable static serving
	if pattern == "/" {
		s.Lock()
		s.static = false
		s.Unlock()
	}

	// register the handler
	s.mux.Handle(pattern, handler)
}

func (s *service) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
	var seen bool
	for _, ep := range s.srv.Endpoints {
		if ep.Name == pattern {
			seen = true
			break
		}
	}
	if !seen {
		s.srv.Endpoints = append(s.srv.Endpoints, &registry.Endpoint{
			Name: pattern,
		})
	}

	s.mux.HandleFunc(pattern, handler)
}

func (s *service) Init(opts ...Option) error {
	for _, o := range opts {
		o(&s.opts)
	}

	serviceOpts := []micro.Option{}

	if len(s.opts.Flags) > 0 {
		serviceOpts = append(serviceOpts, micro.Flags(s.opts.Flags...))
	}

	if s.opts.Registry != nil {
		serviceOpts = append(serviceOpts, micro.Registry(s.opts.Registry))
	}

	serviceOpts = append(serviceOpts, micro.Action(func(ctx *cli.Context) {
		if ttl := ctx.Int("register_ttl"); ttl > 0 {
			s.opts.RegisterTTL = time.Duration(ttl) * time.Second
		}

		if interval := ctx.Int("register_interval"); interval > 0 {
			s.opts.RegisterInterval = time.Duration(interval) * time.Second
		}

		if name := ctx.String("server_name"); len(name) > 0 {
			s.opts.Name = name
		}

		if ver := ctx.String("server_version"); len(ver) > 0 {
			s.opts.Version = ver
		}

		if id := ctx.String("server_id"); len(id) > 0 {
			s.opts.Id = id
		}

		if addr := ctx.String("server_address"); len(addr) > 0 {
			s.opts.Address = addr
		}

		if adv := ctx.String("server_advertise"); len(adv) > 0 {
			s.opts.Advertise = adv
		}

		if s.opts.Action != nil {
			s.opts.Action(ctx)
		}
	}))

	s.opts.Service.Init(serviceOpts...)
	srv := s.genSrv()
	srv.Endpoints = s.srv.Endpoints
	s.srv = srv

	return nil
}

func (s *service) Run() error {
	if err := s.start(); err != nil {
		return err
	}

	if err := s.register(); err != nil {
		return err
	}

	// start reg loop
	ex := make(chan bool)
	go s.run(ex)

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT)

	select {
	// wait on kill signal
	case sig := <-ch:
		log.Logf("Received signal %s\n", sig)
	// wait on context cancel
	case <-s.opts.Context.Done():
		log.Logf("Received context shutdown")
	}

	// exit reg loop
	close(ex)

	if err := s.deregister(); err != nil {
		return err
	}

	return s.stop()
}

// Options returns the options for the given service
func (s *service) Options() Options {
	return s.opts
}

func (s *service) listen(network, addr string) (net.Listener, error) {
	var l net.Listener
	var err error

	// TODO: support use of listen options
	if s.opts.Secure || s.opts.TLSConfig != nil {
		config := s.opts.TLSConfig

		fn := func(addr string) (net.Listener, error) {
			if config == nil {
				hosts := []string{addr}

				// check if its a valid host:port
				if host, _, err := net.SplitHostPort(addr); err == nil {
					if len(host) == 0 {
						hosts = maddr.IPs()
					} else {
						hosts = []string{host}
					}
				}

				// generate a certificate
				cert, err := mls.Certificate(hosts...)
				if err != nil {
					return nil, err
				}
				config = &tls.Config{Certificates: []tls.Certificate{cert}}
			}
			return tls.Listen(network, addr, config)
		}

		l, err = mnet.Listen(addr, fn)
	} else {
		fn := func(addr string) (net.Listener, error) {
			return net.Listen(network, addr)
		}

		l, err = mnet.Listen(addr, fn)
	}

	if err != nil {
		return nil, err
	}

	return l, nil
}
<commit_msg>No newlines at end of log messages<commit_after>package web

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/registry"
	maddr "github.com/micro/go-micro/util/addr"
	mhttp "github.com/micro/go-micro/util/http"
	"github.com/micro/go-micro/util/log"
	mnet "github.com/micro/go-micro/util/net"
	mls "github.com/micro/go-micro/util/tls"
)

type service struct {
	opts Options

	mux *http.ServeMux
	srv *registry.Service

	sync.Mutex
	running bool
	static  bool
	exit    chan chan error
}

func newService(opts ...Option) Service {
	options := newOptions(opts...)
	s := &service{
		opts:   options,
		mux:    http.NewServeMux(),
		static: true,
	}
	s.srv = s.genSrv()
	return s
}

func (s *service) genSrv() *registry.Service {
	// default host:port
	parts := strings.Split(s.opts.Address, ":")
	host := strings.Join(parts[:len(parts)-1], ":")
	port, _ := strconv.Atoi(parts[len(parts)-1])

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(s.opts.Advertise) > 0 {
		parts = strings.Split(s.opts.Advertise, ":")

		// we have host:port
		if len(parts) > 1 {
			// set the host
			host = strings.Join(parts[:len(parts)-1], ":")

			// get the port
			if aport, _ := strconv.Atoi(parts[len(parts)-1]); aport > 0 {
				port = aport
			}
		} else {
			host = parts[0]
		}
	}

	addr, err := maddr.Extract(host)
	if err != nil {
		// best effort localhost
		addr = "127.0.0.1"
	}

	return &registry.Service{
		Name:    s.opts.Name,
		Version: s.opts.Version,
		Nodes: []*registry.Node{{
			Id:       s.opts.Id,
			Address:  fmt.Sprintf("%s:%d", addr, port),
			Metadata: s.opts.Metadata,
		}},
	}
}

func (s *service) run(exit chan bool) {
	if s.opts.RegisterInterval <= time.Duration(0) {
		return
	}

	t := time.NewTicker(s.opts.RegisterInterval)

	for {
		select {
		case <-t.C:
			s.register()
		case <-exit:
			t.Stop()
			return
		}
	}
}

func (s *service) register() error {
	if s.srv == nil {
		return nil
	}
	// default to service registry
	r := s.opts.Service.Client().Options().Registry
	// switch to option if specified
	if s.opts.Registry != nil {
		r = s.opts.Registry
	}

	// service node need modify, node address maybe changed
	srv := s.genSrv()
	srv.Endpoints = s.srv.Endpoints
	s.srv = srv
	return r.Register(s.srv, registry.RegisterTTL(s.opts.RegisterTTL))
}

func (s *service) deregister() error {
	if s.srv == nil {
		return nil
	}
	// default to service registry
	r := s.opts.Service.Client().Options().Registry
	// switch to option if specified
	if s.opts.Registry != nil {
		r = s.opts.Registry
	}
	return r.Deregister(s.srv)
}

func (s *service) start() error {
	s.Lock()
	defer s.Unlock()

	if s.running {
		return nil
	}

	l, err := s.listen("tcp", s.opts.Address)
	if err != nil {
		return err
	}

	s.opts.Address = l.Addr().String()
	srv := s.genSrv()
	srv.Endpoints = s.srv.Endpoints
	s.srv = srv

	var h http.Handler

	if s.opts.Handler != nil {
		h = s.opts.Handler
	} else {
		h = s.mux
		var r sync.Once

		// register the html dir
		r.Do(func() {
			// static dir
			static := s.opts.StaticDir
			if s.opts.StaticDir[0] != '/' {
				dir, _ := os.Getwd()
				static = filepath.Join(dir, static)
			}

			// set static if no / handler is registered
			if s.static {
				_, err := os.Stat(static)
				if err == nil {
					log.Logf("Enabling static file serving from %s", static)
					s.mux.Handle("/", http.FileServer(http.Dir(static)))
				}
			}
		})
	}

	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	var httpSrv *http.Server
	if s.opts.Server != nil {
		httpSrv = s.opts.Server
	} else {
		httpSrv = &http.Server{}
	}

	httpSrv.Handler = h

	go httpSrv.Serve(l)

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	s.exit = make(chan chan error, 1)
	s.running = true

	go func() {
		ch := <-s.exit
		ch <- l.Close()
	}()

	log.Logf("Listening on %v", l.Addr().String())
	return nil
}

func (s *service) stop() error {
	s.Lock()
	defer s.Unlock()

	if !s.running {
		return nil
	}

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			return err
		}
	}

	ch := make(chan error, 1)
	s.exit <- ch
	s.running = false

	log.Log("Stopping")

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			if chErr := <-ch; chErr != nil {
				return chErr
			}
			return err
		}
	}

	return <-ch
}

func (s *service) Client() *http.Client {
	rt := mhttp.NewRoundTripper(
		mhttp.WithRegistry(registry.DefaultRegistry),
	)
	return &http.Client{
		Transport: rt,
	}
}

func (s *service) Handle(pattern string, handler http.Handler) {
	var seen bool
	for _, ep := range s.srv.Endpoints {
		if ep.Name == pattern {
			seen = true
			break
		}
	}

	// if its unseen then add an endpoint
	if !seen {
		s.srv.Endpoints = append(s.srv.Endpoints, &registry.Endpoint{
			Name: pattern,
		})
	}

	// disable static serving
	if pattern == "/" {
		s.Lock()
		s.static = false
		s.Unlock()
	}

	// register the handler
	s.mux.Handle(pattern, handler)
}

func (s *service) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
	var seen bool
	for _, ep := range s.srv.Endpoints {
		if ep.Name == pattern {
			seen = true
			break
		}
	}
	if !seen {
		s.srv.Endpoints = append(s.srv.Endpoints, &registry.Endpoint{
			Name: pattern,
		})
	}

	s.mux.HandleFunc(pattern, handler)
}

func (s *service) Init(opts ...Option) error {
	for _, o := range opts {
		o(&s.opts)
	}

	serviceOpts := []micro.Option{}

	if len(s.opts.Flags) > 0 {
		serviceOpts = append(serviceOpts, micro.Flags(s.opts.Flags...))
	}

	if s.opts.Registry != nil {
		serviceOpts = append(serviceOpts, micro.Registry(s.opts.Registry))
	}

	serviceOpts = append(serviceOpts, micro.Action(func(ctx *cli.Context) {
		if ttl := ctx.Int("register_ttl"); ttl > 0 {
			s.opts.RegisterTTL = time.Duration(ttl) * time.Second
		}

		if interval := ctx.Int("register_interval"); interval > 0 {
			s.opts.RegisterInterval = time.Duration(interval) * time.Second
		}

		if name := ctx.String("server_name"); len(name) > 0 {
			s.opts.Name = name
		}

		if ver := ctx.String("server_version"); len(ver) > 0 {
			s.opts.Version = ver
		}

		if id := ctx.String("server_id"); len(id) > 0 {
			s.opts.Id = id
		}

		if addr := ctx.String("server_address"); len(addr) > 0 {
			s.opts.Address = addr
		}

		if adv := ctx.String("server_advertise"); len(adv) > 0 {
			s.opts.Advertise = adv
		}

		if s.opts.Action != nil {
			s.opts.Action(ctx)
		}
	}))

	s.opts.Service.Init(serviceOpts...)
	srv := s.genSrv()
	srv.Endpoints = s.srv.Endpoints
	s.srv = srv

	return nil
}

func (s *service) Run() error {
	if err := s.start(); err != nil {
		return err
	}

	if err := s.register(); err != nil {
		return err
	}

	// start reg loop
	ex := make(chan bool)
	go s.run(ex)

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT)

	select {
	// wait on kill signal
	case sig := <-ch:
		log.Logf("Received signal %s", sig)
	// wait on context cancel
	case <-s.opts.Context.Done():
		log.Logf("Received context shutdown")
	}

	// exit reg loop
	close(ex)

	if err := s.deregister(); err != nil {
		return err
	}

	return s.stop()
}

// Options returns the options for the given service
func (s *service) Options() Options {
	return s.opts
}

func (s *service) listen(network, addr string) (net.Listener, error) {
	var l net.Listener
	var err error

	// TODO: support use of listen options
	if s.opts.Secure || s.opts.TLSConfig != nil {
		config := s.opts.TLSConfig

		fn := func(addr string) (net.Listener, error) {
			if config == nil {
				hosts := []string{addr}

				// check if its a valid host:port
				if host, _, err := net.SplitHostPort(addr); err == nil {
					if len(host) == 0 {
						hosts = maddr.IPs()
					} else {
						hosts = []string{host}
					}
				}

				// generate a certificate
				cert, err := mls.Certificate(hosts...)
				if err != nil {
					return nil, err
				}
				config = &tls.Config{Certificates: []tls.Certificate{cert}}
			}
			return tls.Listen(network, addr, config)
		}

		l, err = mnet.Listen(addr, fn)
	} else {
		fn := func(addr string) (net.Listener, error) {
			return net.Listen(network, addr)
		}

		l, err = mnet.Listen(addr, fn)
	}

	if err != nil {
		return nil, err
	}

	return l, nil
}
<|endoftext|>


<commit_before>package topbar

import (
	"github.com/ambientsound/pms/api"
	"github.com/ambientsound/pms/mpd"
)

var stateStrings = map[string]string{
	mpd.StatePlay:    "|>",
	mpd.StatePause:   "||",
	mpd.StateStop:    "[]",
	mpd.StateUnknown: "??",
}

// State draws the current player state as an ASCII symbol.
type State struct {
	api api.API
}

// NewState returns State.
func NewState(a api.API, param string) Fragment {
	return &State{a}
}

// Text implements Fragment.
func (w *State) Text() (string, string) {
	playerStatus := w.api.PlayerStatus()
	return stateStrings[playerStatus.State], `state`
}
<commit_msg>Offer unicode strings in topbar fragment ${state|unicode}<commit_after>package topbar

import (
	"github.com/ambientsound/pms/api"
	"github.com/ambientsound/pms/mpd"
)

var stateStrings = map[string]string{
	mpd.StatePlay:    "|>",
	mpd.StatePause:   "||",
	mpd.StateStop:    "[]",
	mpd.StateUnknown: "??",
}

var stateUnicodes = map[string]string{
	mpd.StatePlay:    "\u25b6",
	mpd.StatePause:   "\u23f8",
	mpd.StateStop:    "\u23f9",
	mpd.StateUnknown: "\u2bd1",
}

// State draws the current player state as an ASCII symbol.
type State struct {
	api   api.API
	table map[string]string
}

// NewState returns State.
func NewState(a api.API, param string) Fragment {
	table := stateStrings
	if param == "unicode" {
		table = stateUnicodes
	}
	return &State{a, table}
}

// Text implements Fragment.
func (w *State) Text() (string, string) {
	playerStatus := w.api.PlayerStatus()
	return w.table[playerStatus.State], `state`
}
<|endoftext|>


<commit_before>package script

import (
	"errors"
	"fmt"
	"github.com/aarzilli/golua/lua"
	anko_parser "github.com/mattn/anko/parser"
	anko "github.com/mattn/anko/vm"
	"github.com/robertkrimen/otto"
	"github.com/tyler-sommer/squircy2/squircy/event"
	"github.com/tyler-sommer/squircy2/squircy/irc"
	glisp "github.com/zhemao/glisp/interpreter"
	"log"
	"reflect"
	"time"
)

const maxExecutionTime = 2 // in seconds
var Halt = errors.New("Execution limit exceeded")
var UnknownScriptType = errors.New("Unknown script type")

type ScriptType string

const (
	Javascript ScriptType = "Javascript"
	Lua                   = "Lua"
	Lisp                  = "Lisp"
	Anko                  = "Anko"
)

type ScriptManager struct {
	e            event.EventManager
	jsVm         *otto.Otto
	jsDriver     javascriptDriver
	luaVm        *lua.State
	luaDriver    luaDriver
	lispVm       *glisp.Glisp
	lispDriver   lispDriver
	ankoVm       *anko.Env
	ankoDriver   ankoDriver
	httpHelper   httpHelper
	ircHelper    ircHelper
	dataHelper   dataHelper
	scriptHelper scriptHelper
	repo         ScriptRepository
	l            *log.Logger
}

func NewScriptManager(repo ScriptRepository, l *log.Logger, e event.EventManager) ScriptManager {
	mgr := ScriptManager{
		e,
		nil,
		javascriptDriver{},
		nil,
		luaDriver{},
		nil,
		lispDriver{},
		nil,
		ankoDriver{},
		httpHelper{},
		ircHelper{},
		dataHelper{make(map[string]interface{})},
		scriptHelper{},
		repo,
		l,
	}
	mgr.init()

	return mgr
}

func (m *ScriptManager) RunUnsafe(t ScriptType, code string) (result interface{}, err error) {
	err = nil
	result = nil

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

	switch {
	case t == Javascript:
		res, e := runUnsafeJavascript(m.jsVm, code)
		if e != nil {
			err = e
			return
		}
		result, _ = res.Export()

	case t == Lua:
		err = runUnsafeLua(m.luaVm, code)

	case t == Lisp:
		res, e := runUnsafeLisp(m.lispVm, code)
		if e != nil {
			err = e
			return
		}
		result = res.SexpString()

	case t == Anko:
		res, e := runUnsafeAnko(m.ankoVm, code)
		if e != nil {
			err = e
			return
		}
		result = res.Interface()

	default:
		err = UnknownScriptType
	}

	return
}

func runUnsafeJavascript(vm *otto.Otto, unsafe string) (val otto.Value, err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e == Halt {
				fmt.Println("Some code took too long! Stopping after: ", duration)
			}
			err = e.(error)
		}
	}()

	vm.Interrupt = make(chan func(), 1)

	go func() {
		time.Sleep(maxExecutionTime * time.Second)
		vm.Interrupt <- func() {
			panic(Halt)
		}
	}()

	val, err = vm.Run(unsafe)

	return
}

func runUnsafeLua(vm *lua.State, unsafe string) (err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e == Halt {
				fmt.Println("Some code took too long! Stopping after: ", duration)
			}
			err = e.(error)
		}
	}()

	vm.SetExecutionLimit(maxExecutionTime * (1 << 26))
	err = vm.DoString(unsafe)

	return
}

func runUnsafeLisp(vm *glisp.Glisp, unsafe string) (val glisp.Sexp, err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e.(error).Error() == "Execution limit exceeded" {
				fmt.Println("Some code took too long! Stopping after: ", duration)
				e = Halt
			}
			val = glisp.SexpNull
			err = e.(error)
		}
	}()

	vm.Clear()
	vm.LoadString(unsafe)
	val, err = vm.Run()

	return
}

func runUnsafeAnko(vm *anko.Env, unsafe string) (val reflect.Value, err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e == Halt {
				fmt.Println("Some code took too long! Stopping after: ", duration)
			}
			val = anko.NilValue
			err = e.(error)
		}
	}()

	scanner := &anko_parser.Scanner{}
	scanner.Init(unsafe)
	stmts, err := anko_parser.Parse(scanner)
	if err != nil {
		val = anko.NilValue
		return
	}
	val, err = anko.Run(stmts, vm)

	return
}

func (m *ScriptManager) ReInit() {
	m.init()
}

func (m *ScriptManager) init() {
	m.e.ClearAll()
	m.e.Bind(irc.ConnectingEvent, func(mgr *irc.IrcConnectionManager) {
		m.ircHelper.conn = mgr.Connection()
	})

	m.e.Bind(irc.DisconnectEvent, func(ev event.Event) {
		m.ircHelper.conn = nil
	})

	jsVm := newJavascriptVm(m)
	luaVm := newLuaVm(m)
	lispVm := newLispVm(m)
	ankoVm := newAnkoVm(m)

	m.jsVm = jsVm
	m.jsDriver.vm = jsVm
	m.luaVm = luaVm
	m.luaDriver.vm = luaVm
	m.lispVm = lispVm
	m.lispDriver.vm = lispVm
	m.ankoVm = ankoVm
	m.ankoDriver.vm = ankoVm

	m.scriptHelper = scriptHelper{m.e, m.jsDriver, m.luaDriver, m.lispDriver, m.ankoDriver, make(map[string]event.EventHandler, 0)}

	scripts := m.repo.FetchAll()
	for _, script := range scripts {
		m.l.Println("Running", script.Type, "script", script.Title)
		m.RunUnsafe(script.Type, script.Body)
	}
}
<commit_msg>Added rudimentary Sexp conversion function<commit_after>package script

import (
	"errors"
	"fmt"
	"github.com/aarzilli/golua/lua"
	anko_parser "github.com/mattn/anko/parser"
	anko "github.com/mattn/anko/vm"
	"github.com/robertkrimen/otto"
	"github.com/tyler-sommer/squircy2/squircy/event"
	"github.com/tyler-sommer/squircy2/squircy/irc"
	glisp "github.com/zhemao/glisp/interpreter"
	"log"
	"reflect"
	"time"
)

const maxExecutionTime = 2 // in seconds
var Halt = errors.New("Execution limit exceeded")
var UnknownScriptType = errors.New("Unknown script type")

type ScriptType string

const (
	Javascript ScriptType = "Javascript"
	Lua                   = "Lua"
	Lisp                  = "Lisp"
	Anko                  = "Anko"
)

type ScriptManager struct {
	e            event.EventManager
	jsVm         *otto.Otto
	jsDriver     javascriptDriver
	luaVm        *lua.State
	luaDriver    luaDriver
	lispVm       *glisp.Glisp
	lispDriver   lispDriver
	ankoVm       *anko.Env
	ankoDriver   ankoDriver
	httpHelper   httpHelper
	ircHelper    ircHelper
	dataHelper   dataHelper
	scriptHelper scriptHelper
	repo         ScriptRepository
	l            *log.Logger
}

func NewScriptManager(repo ScriptRepository, l *log.Logger, e event.EventManager) ScriptManager {
	mgr := ScriptManager{
		e,
		nil,
		javascriptDriver{},
		nil,
		luaDriver{},
		nil,
		lispDriver{},
		nil,
		ankoDriver{},
		httpHelper{},
		ircHelper{},
		dataHelper{make(map[string]interface{})},
		scriptHelper{},
		repo,
		l,
	}
	mgr.init()

	return mgr
}

func (m *ScriptManager) RunUnsafe(t ScriptType, code string) (result interface{}, err error) {
	err = nil
	result = nil

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

	switch {
	case t == Javascript:
		res, e := runUnsafeJavascript(m.jsVm, code)
		if e != nil {
			err = e
			return
		}
		result, _ = res.Export()

	case t == Lua:
		err = runUnsafeLua(m.luaVm, code)

	case t == Lisp:
		res, e := runUnsafeLisp(m.lispVm, code)
		if e != nil {
			err = e
			return
		}
		result = sexpToInterface(res)

	case t == Anko:
		res, e := runUnsafeAnko(m.ankoVm, code)
		if e != nil {
			err = e
			return
		}
		result = res.Interface()

	default:
		err = UnknownScriptType
	}

	return
}

func runUnsafeJavascript(vm *otto.Otto, unsafe string) (val otto.Value, err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e == Halt {
				fmt.Println("Some code took too long! Stopping after: ", duration)
			}
			err = e.(error)
		}
	}()

	vm.Interrupt = make(chan func(), 1)

	go func() {
		time.Sleep(maxExecutionTime * time.Second)
		vm.Interrupt <- func() {
			panic(Halt)
		}
	}()

	val, err = vm.Run(unsafe)

	return
}

func runUnsafeLua(vm *lua.State, unsafe string) (err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e == Halt {
				fmt.Println("Some code took too long! Stopping after: ", duration)
			}
			err = e.(error)
		}
	}()

	vm.SetExecutionLimit(maxExecutionTime * (1 << 26))
	err = vm.DoString(unsafe)

	return
}

func runUnsafeLisp(vm *glisp.Glisp, unsafe string) (val glisp.Sexp, err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e.(error).Error() == "Execution limit exceeded" {
				fmt.Println("Some code took too long! Stopping after: ", duration)
				e = Halt
			}
			val = glisp.SexpNull
			err = e.(error)
		}
	}()

	vm.Clear()
	vm.LoadString(unsafe)
	val, err = vm.Run()

	return
}

func sexpToInterface(val glisp.Sexp) interface{} {
	switch t := val.(type) {
	case glisp.SexpInt:
		return int(t)

	case glisp.SexpFloat:
		return float64(t)

	case glisp.SexpStr:
		return string(t)

	default:
		return val.SexpString()
	}
}

func runUnsafeAnko(vm *anko.Env, unsafe string) (val reflect.Value, err error) {
	start := time.Now()
	defer func() {
		duration := time.Since(start)
		if e := recover(); e != nil {
			if e == Halt {
				fmt.Println("Some code took too long! Stopping after: ", duration)
			}
			val = anko.NilValue
			err = e.(error)
		}
	}()

	scanner := &anko_parser.Scanner{}
	scanner.Init(unsafe)
	stmts, err := anko_parser.Parse(scanner)
	if err != nil {
		val = anko.NilValue
		return
	}
	val, err = anko.Run(stmts, vm)

	return
}

func (m *ScriptManager) ReInit() {
	m.init()
}

func (m *ScriptManager) init() {
	m.e.ClearAll()
	m.e.Bind(irc.ConnectingEvent, func(mgr *irc.IrcConnectionManager) {
		m.ircHelper.conn = mgr.Connection()
	})

	m.e.Bind(irc.DisconnectEvent, func(ev event.Event) {
		m.ircHelper.conn = nil
	})

	jsVm := newJavascriptVm(m)
	luaVm := newLuaVm(m)
	lispVm := newLispVm(m)
	ankoVm := newAnkoVm(m)

	m.jsVm = jsVm
	m.jsDriver.vm = jsVm
	m.luaVm = luaVm
	m.luaDriver.vm = luaVm
	m.lispVm = lispVm
	m.lispDriver.vm = lispVm
	m.ankoVm = ankoVm
	m.ankoDriver.vm = ankoVm

	m.scriptHelper = scriptHelper{m.e, m.jsDriver, m.luaDriver, m.lispDriver, m.ankoDriver, make(map[string]event.EventHandler, 0)}

	scripts := m.repo.FetchAll()
	for _, script := range scripts {
		m.l.Println("Running", script.Type, "script", script.Title)
		m.RunUnsafe(script.Type, script.Body)
	}
}
<|endoftext|>