<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/log"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

func main() {
	lifecycle.Startup("broker", false)
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := log.CreateCounterGauge("newClients", log.NoUnit, true)
	changeWebsocketClientsGauge := log.CreateCounterGauge("websocketClients", log.NoUnit, false)
	log.RunGaugesLoop()

	publishConn := amqputil.CreateConnection("broker")
	defer publishConn.Close()

	routeMap := make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap := make(map[string]*map[string]bool)
	var globalMapMutex sync.Mutex

	service := sockjs.NewService(config.Current.Client.StaticFilesBaseUrl+"/js/sock.js", 10*time.Minute, func(session *sockjs.Session) {
		defer log.RecoverAndLog()

		r := make([]byte, 128/8)
		rand.Read(r)
		socketId := base64.StdEncoding.EncodeToString(r)
		session.Tag = socketId

		log.Debug("Client connected: " + socketId)
		changeClientsGauge(1)
		changeNewClientsGauge(1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(1)
		}
		defer func() {
			log.Debug("Client disconnected: " + socketId)
			changeClientsGauge(-1)
			if session.IsWebsocket {
				changeWebsocketClientsGauge(-1)
			}
		}()

		var controlChannel *amqp.Channel
		var lastPayload string
		resetControlChannel := func() {
			if controlChannel != nil {
				controlChannel.Close()
			}
			var err error
			controlChannel, err = publishConn.Channel()
			if err != nil {
				panic(err)
			}
			go func() {
				defer log.RecoverAndLog()

				for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
					log.Warn("AMQP channel: "+amqpErr.Error(), "Last publish payload:", lastPayload)

					sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
				}
			}()
		}
		resetControlChannel()
		defer func() { controlChannel.Close() }()

		subscriptions := make(map[string]bool)
		globalMapMutex.Lock()
		socketSubscriptionsMap[socketId] = &subscriptions
		globalMapMutex.Unlock()

		removeFromRouteMap := func(routingKeyPrefix string) {
			routeSessions := routeMap[routingKeyPrefix]
			for i, routeSession := range routeSessions {
				if routeSession == session {
					routeSessions[i] = routeSessions[len(routeSessions)-1]
					routeSessions = routeSessions[:len(routeSessions)-1]
					break
				}
			}
			if len(routeSessions) == 0 {
				delete(routeMap, routingKeyPrefix)
				return
			}
			routeMap[routingKeyPrefix] = routeSessions
		}

		subscribe := func(routingKeyPrefix string) {
			if subscriptions[routingKeyPrefix] {
				log.Warn("Duplicate subscription to same routing key.", session.Tag, routingKeyPrefix)
				return
			}
			if len(subscriptions) > 0 && len(subscriptions)%1000 == 0 {
				log.Warn("Client with more than "+strconv.Itoa(len(subscriptions))+" subscriptions.", session.Tag)
			}
			routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
			subscriptions[routingKeyPrefix] = true
		}

		unsubscribe := func(routingKeyPrefix string) {
			removeFromRouteMap(routingKeyPrefix)
			delete(subscriptions, routingKeyPrefix)
		}

		defer func() {
			globalMapMutex.Lock()
			for routingKeyPrefix := range subscriptions {
				removeFromRouteMap(routingKeyPrefix)
			}
			globalMapMutex.Unlock()

			time.AfterFunc(5*time.Minute, func() {
				globalMapMutex.Lock()
				delete(socketSubscriptionsMap, socketId)
				globalMapMutex.Unlock()
			})

			for {
				err := controlChannel.Publish("authAll", "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
				if err == nil {
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					panic(err)
				}
				resetControlChannel()
			}
		}()

		err := controlChannel.Publish("authAll", "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
		if err != nil {
			panic(err)
		}

		sendToClient(session, "broker.connected", socketId)

		for data := range session.ReceiveChan {
			if data == nil || session.Closed {
				break
			}
			func() {
				defer log.RecoverAndLog()

				message := data.(map[string]interface{})
				log.Debug("Received message", message)

				action := message["action"]
				switch action {
				case "subscribe":
					globalMapMutex.Lock()
					defer globalMapMutex.Unlock()
					for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
						subscribe(routingKeyPrefix)
					}
					sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

				case "resubscribe":
					globalMapMutex.Lock()
					defer globalMapMutex.Unlock()
					oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
					if found {
						for routingKeyPrefix := range *oldSubscriptions {
							subscribe(routingKeyPrefix)
						}<commit_msg>broker: Receive client version with "clientInfo" action.<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/log"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

func main() {
	lifecycle.Startup("broker", false)
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := log.CreateCounterGauge("newClients", log.NoUnit, true)
	changeWebsocketClientsGauge := log.CreateCounterGauge("websocketClients", log.NoUnit, false)
	log.RunGaugesLoop()

	publishConn := amqputil.CreateConnection("broker")
	defer publishConn.Close()

	routeMap := make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap := make(map[string]*map[string]bool)
	var globalMapMutex sync.Mutex

	service := sockjs.NewService(config.Current.Client.StaticFilesBaseUrl+"/js/sock.js", 10*time.Minute, func(session *sockjs.Session) {
		defer log.RecoverAndLog()

		r := make([]byte, 128/8)
		rand.Read(r)
		socketId := base64.StdEncoding.EncodeToString(r)
		session.Tag = socketId
		clientVersion := 0

		log.Debug("Client connected: " + socketId)
		changeClientsGauge(1)
		changeNewClientsGauge(1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(1)
		}
		defer func() {
			log.Debug("Client disconnected: " + socketId)
			changeClientsGauge(-1)
			if session.IsWebsocket {
				changeWebsocketClientsGauge(-1)
			}
		}()

		var controlChannel *amqp.Channel
		var lastPayload string
		resetControlChannel := func() {
			if controlChannel != nil {
				controlChannel.Close()
			}
			var err error
			controlChannel, err = publishConn.Channel()
			if err != nil {
				panic(err)
			}
			go func() {
				defer log.RecoverAndLog()

				for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
					log.Warn("AMQP channel: "+amqpErr.Error(), "Last publish payload:", lastPayload)

					sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
				}
			}()
		}
		resetControlChannel()
		defer func() { controlChannel.Close() }()

		subscriptions := make(map[string]bool)
		globalMapMutex.Lock()
		socketSubscriptionsMap[socketId] = &subscriptions
		globalMapMutex.Unlock()

		removeFromRouteMap := func(routingKeyPrefix string) {
			routeSessions := routeMap[routingKeyPrefix]
			for i, routeSession := range routeSessions {
				if routeSession == session {
					routeSessions[i] = routeSessions[len(routeSessions)-1]
					routeSessions = routeSessions[:len(routeSessions)-1]
					break
				}
			}
			if len(routeSessions) == 0 {
				delete(routeMap, routingKeyPrefix)
				return
			}
			routeMap[routingKeyPrefix] = routeSessions
		}

		subscribe := func(routingKeyPrefix string) {
			if subscriptions[routingKeyPrefix] {
				log.Warn("Duplicate subscription to same routing key.", session.Tag, routingKeyPrefix)
				return
			}
			if len(subscriptions) > 0 && len(subscriptions)%1000 == 0 {
				log.Warn("Client with more than "+strconv.Itoa(len(subscriptions))+" subscriptions.", session.Tag)
			}
			routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
			subscriptions[routingKeyPrefix] = true
		}

		unsubscribe := func(routingKeyPrefix string) {
			removeFromRouteMap(routingKeyPrefix)
			delete(subscriptions, routingKeyPrefix)
		}

		defer func() {
			globalMapMutex.Lock()
			for routingKeyPrefix := range subscriptions {
				removeFromRouteMap(routingKeyPrefix)
			}
			globalMapMutex.Unlock()

			time.AfterFunc(5*time.Minute, func() {
				globalMapMutex.Lock()
				delete(socketSubscriptionsMap, socketId)
				globalMapMutex.Unlock()
			})

			for {
				err := controlChannel.Publish("authAll", "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
				if err == nil {
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					panic(err)
				}
				resetControlChannel()
			}
		}()

		err := controlChannel.Publish("authAll", "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
		if err != nil {
			panic(err)
		}

		sendToClient(session, "broker.connected", socketId)

		for data := range session.ReceiveChan {
			if data == nil || session.Closed {
				break
			}
			func() {
				defer log.RecoverAndLog()

				message := data.(map[string]interface{})
				log.Debug("Received message", message)

				action := message["action"]
				switch action {
				case "clientInfo":
					clientVersion = message["version"].(int)

				case "subscribe":
					globalMapMutex.Lock()
					defer globalMapMutex.Unlock()
					for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
						subscribe(routingKeyPrefix)
					}
					sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

				case "resubscribe":
					globalMapMutex.Lock()
					defer globalMapMutex.Unlock()
					oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
					if found {
						for routingKeyPrefix := range *oldSubscriptions {
							subscribe(routingKeyPrefix)
						}
{'fraction_non_alphanumeric': 0.13166351606805293, 'fraction_numerical': 0.005387523629489603, 'mean_word_length': 4.214884179398719, 'pattern_counts': {'":': 13, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, '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': '752341', 'n_tokens_mistral': 4116, 'n_tokens_neox': 3476, 'n_words': 618}