<commit_before>package kuja

import (
	"crypto/tls"
	log "github.com/Sirupsen/logrus"
	"github.com/golang/snappy/snappy"
	"github.com/plimble/kuja/broker"
	"github.com/plimble/kuja/encoder"
	"github.com/plimble/kuja/encoder/json"
	"github.com/plimble/kuja/registry"
	"github.com/satori/go.uuid"
	"golang.org/x/net/netutil"
	"gopkg.in/tylerb/graceful.v1"
	"io"
	"net"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"
)

type Handler func(ctx *Context, w http.ResponseWriter, r *http.Request) error

type Server struct {
	id              string
	pool            sync.Pool
	middleware      []Handler
	mu              sync.Mutex // protects the serviceMap
	serviceMap      map[string]*service
	subscriberMap   map[string]*subscriber
	broker          broker.Broker
	encoder         encoder.Encoder
	snappy          bool
	serviceError    ServiceErrorFunc
	subscriberError SubscriberErrorFunc
	registry        registry.Registry
	srv             *graceful.Server
}

func NewServer() *Server {
	server := &Server{
		id:              uuid.NewV1().String(),
		serviceMap:      make(map[string]*service),
		encoder:         json.NewEncoder(),
		serviceError:    defaulServiceErr,
		subscriberError: defaulSubscriberErr,
		subscriberMap:   make(map[string]*subscriber),
	}

	server.pool.New = func() interface{} {
		return &Context{
			ReqMetadata:  make(Metadata),
			RespMetadata: make(Metadata),
			returnValues: make([]reflect.Value, 1),
			serviceError: server.serviceError,
			isResp:       false,
		}
	}

	return server
}

func (server *Server) Use(h ...Handler) {
	server.middleware = append(server.middleware, h...)
}

func (server *Server) Snappy(enable bool) {
	server.snappy = enable
}

func (server *Server) Registry(r registry.Registry) {
	server.registry = r
}

func (server *Server) Encoder(enc encoder.Encoder) {
	server.encoder = enc
}

func (server *Server) Broker(b broker.Broker) {
	server.broker = b
}

func (server *Server) Run(addr string, timeout time.Duration) {
	server.srv = &graceful.Server{
		Timeout: timeout,
		Server: &http.Server{
			Addr:    addr,
			Handler: server,
		},
	}

	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error(err)
		return
	}

	if server.srv.ListenLimit != 0 {
		l = netutil.LimitListener(l, server.srv.ListenLimit)
	}

	if err := server.startRegistry("http://", addr); err != nil {
		log.Error(err)
		return
	}

	if err := server.startBroker(); err != nil {
		log.Error(err)
		return
	}

	if err := server.startSubscribe(); err != nil {
		log.Error(err)
		return
	}

	log.Infof("Start server id %s on %s", server.id, addr)
	server.srv.Serve(l)
	log.Info("Stop server")
	server.stopRegistry()
	server.stopBroker()
}

func (server *Server) Close() {
	server.srv.Stop(0)
}

func (server *Server) RunTLS(addr string, timeout time.Duration, certFile, keyFile string) {
	server.srv = &graceful.Server{
		Timeout: timeout,
		Server: &http.Server{
			Addr:    addr,
			Handler: server,
		},
	}

	config := &tls.Config{}
	if server.srv.TLSConfig != nil {
		*config = *server.srv.TLSConfig
	}

	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Error(err)
		return
	}

	conn, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error(err)
		return
	}

	tlsListener := tls.NewListener(conn, config)

	if err := server.startRegistry("https://", addr); err != nil {
		log.Error(err)
		return
	}

	if err := server.startBroker(); err != nil {
		log.Error(err)
		return
	}

	if err := server.startSubscribe(); err != nil {
		log.Error(err)
		return
	}

	log.Infof("Start server id %s on %s", server.id, addr)
	server.srv.Serve(tlsListener)
	log.Info("Stop server")
	server.stopRegistry()
	server.stopBroker()
}

func (server *Server) startRegistry(scheme, addr string) error {
	if server.registry == nil {
		return nil
	}

	var host string
	var port string
	parts := strings.Split(addr, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port = parts[len(parts)-1]
	} else {
		host = parts[0]
	}
<commit_msg>add stage<commit_after>package kuja

import (
	"crypto/tls"
	log "github.com/Sirupsen/logrus"
	"github.com/golang/snappy/snappy"
	"github.com/plimble/kuja/broker"
	"github.com/plimble/kuja/encoder"
	"github.com/plimble/kuja/encoder/json"
	"github.com/plimble/kuja/registry"
	"github.com/satori/go.uuid"
	"golang.org/x/net/netutil"
	"gopkg.in/tylerb/graceful.v1"
	"io"
	"net"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"
)

const (
	DEVELOPMENT = "development"
	STAGING     = "staging"
	PRODUCTION  = "production"
)

type Handler func(ctx *Context, w http.ResponseWriter, r *http.Request) error

type Server struct {
	stage           string
	id              string
	pool            sync.Pool
	middleware      []Handler
	mu              sync.Mutex // protects the serviceMap
	serviceMap      map[string]*service
	subscriberMap   map[string]*subscriber
	broker          broker.Broker
	encoder         encoder.Encoder
	snappy          bool
	serviceError    ServiceErrorFunc
	subscriberError SubscriberErrorFunc
	registry        registry.Registry
	srv             *graceful.Server
}

func NewServer() *Server {
	server := &Server{
		stage:           DEVELOPMENT,
		id:              uuid.NewV1().String(),
		serviceMap:      make(map[string]*service),
		encoder:         json.NewEncoder(),
		serviceError:    defaulServiceErr,
		subscriberError: defaulSubscriberErr,
		subscriberMap:   make(map[string]*subscriber),
	}

	server.pool.New = func() interface{} {
		return &Context{
			ReqMetadata:  make(Metadata),
			RespMetadata: make(Metadata),
			returnValues: make([]reflect.Value, 1),
			serviceError: server.serviceError,
			isResp:       false,
		}
	}

	return server
}

func (server *Server) Use(h ...Handler) {
	server.middleware = append(server.middleware, h...)
}

func (server *Server) Stage(stage string) {
	server.stage = stage
}

func (server *Server) Snappy(enable bool) {
	server.snappy = enable
}

func (server *Server) Registry(r registry.Registry) {
	server.registry = r
}

func (server *Server) Encoder(enc encoder.Encoder) {
	server.encoder = enc
}

func (server *Server) Broker(b broker.Broker) {
	server.broker = b
}

func (server *Server) Run(addr string, timeout time.Duration) {
	server.srv = &graceful.Server{
		Timeout: timeout,
		Server: &http.Server{
			Addr:    addr,
			Handler: server,
		},
	}

	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error(err)
		return
	}

	if server.srv.ListenLimit != 0 {
		l = netutil.LimitListener(l, server.srv.ListenLimit)
	}

	if err := server.startRegistry("http://", addr); err != nil {
		log.Error(err)
		return
	}

	if err := server.startBroker(); err != nil {
		log.Error(err)
		return
	}

	if err := server.startSubscribe(); err != nil {
		log.Error(err)
		return
	}

	log.Infof("Start server id %s on %s stage %s", server.id, addr, server.stage)
	server.srv.Serve(l)
	log.Info("Stop server")
	server.stopRegistry()
	server.stopBroker()
}

func (server *Server) Close() {
	server.srv.Stop(0)
}

func (server *Server) RunTLS(addr string, timeout time.Duration, certFile, keyFile string) {
	server.srv = &graceful.Server{
		Timeout: timeout,
		Server: &http.Server{
			Addr:    addr,
			Handler: server,
		},
	}

	config := &tls.Config{}
	if server.srv.TLSConfig != nil {
		*config = *server.srv.TLSConfig
	}

	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Error(err)
		return
	}

	conn, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error(err)
		return
	}

	tlsListener := tls.NewListener(conn, config)

	if err := server.startRegistry("https://", addr); err != nil {
		log.Error(err)
		return
	}

	if err := server.startBroker(); err != nil {
		log.Error(err)
		return
	}

	if err := server.startSubscribe(); err != nil {
		log.Error(err)
		return
	}

	log.Infof("Start server id %s on %s stage %s", server.id, addr, server.stage)
	server.srv.Serve(tlsListener)
	log.Info("Stop server")
	server.stopRegistry()
	server.stopBroker()
}

func (server *Server) startRegistry(scheme, addr string) error {
	if server.registry == nil {
		return nil
	}

	var host string
	var port string
	parts := strings.Split(addr, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port = parts[len(parts)-1]
	} else {
		host = parts[0]
	}

{'fraction_non_alphanumeric': 0.14328984156570362, 'fraction_numerical': 0.003727865796831314, 'mean_word_length': 3.377868434472208, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3155665', 'n_tokens_mistral': 3279, 'n_tokens_neox': 3089, 'n_words': 683}