<commit_before>package template

import (
	. "github.com/dave/jennifer/jen"
	"github.com/devimteam/microgen/parser"
	"github.com/devimteam/microgen/util"
)

var (
	defaultProtoTypes = []string{"string", "bool", "byte", "int64", "uint64", "float64", "int32", "uint32", "float32"}
	goToProtoTypesMap = map[string]string{
		"uint": "uint64",
		"int":  "int64",
	}
	defaultGolangTypes = []string{"string", "bool", "int", "uint", "byte", "int64", "uint64", "float64", "int32", "uint32", "float32"}
)

type GRPCConverterTemplate struct {
	PackagePath string
}

func utilPackagePath(path string) string {
	return path + "/util"
}

func (t GRPCClientTemplate) converterPackagePath() string {
	return t.PackagePath + "/transport/converter/protobuf"
}

// Renders converter file.
//
//		// This file was automatically generated by "microgen" utility.
//		// Please, do not edit.
//		package protobuf
//
//		import (
//			context "context"
//			grpc "github.com/devimteam/go-kit/transportlayer/grpc"
//			svc "github.com/devimteam/microgen/test/svc"
//			util "github.com/devimteam/microgen/test/svc/util"
//			stringsvc "gitlab.devim.team/protobuf/stringsvc"
//		)
//
//		var CountConverter = &grpc.EndpointConverter{
//			func(_ context.Context, request interface{}) (interface{}, error) {
//				req := request.(*svc.CountRequest)
//				return &stringsvc.CountRequest{
//					Symbol: req.Symbol,
//					Text:   req.Text,
//				}, nil
//			},
//			func(_ context.Context, response interface{}) (interface{}, error) {
//				resp := response.(*svc.CountResponse)
//				respPositions, err := util.IntListToProto(resp.Positions)
//				if err != nil {
//					return nil, err
//				}
//				return &stringsvc.CountResponse{
//					Count:     int64(resp.Count),
//					Positions: respPositions,
//				}, nil
//			},
//			func(_ context.Context, request interface{}) (interface{}, error) {
//				req := request.(*stringsvc.CountRequest)
//				return &svc.CountRequest{
//					Symbol: string(req.Symbol),
//					Text:   string(req.Text),
//				}, nil
//			},
//			func(_ context.Context, response interface{}) (interface{}, error) {
//				resp := response.(*stringsvc.CountResponse)
//				respPositions, err := util.ProtoToIntList(resp.Positions)
//				if err != nil {
//					return nil, err
//				}
//				return &svc.CountResponse{
//					Count:     int(resp.Count),
//					Positions: respPositions,
//				}, nil
//			},
//			stringsvc.CountResponse{},
//		}
//
func (t GRPCConverterTemplate) Render(i *parser.Interface) *File {
	f := NewFile("protobuf")

	for _, signature := range i.FuncSignatures {
		f.Var().Id(converterStructName(signature)).Op("=").Op("&").Qual(PackagePathTransportLayerGRPC, "EndpointConverter").
			ValuesFunc(func(g *Group) {
				g.Add(t.encodeRequest(signature, i))
				g.Add(t.encodeResponse(signature, i))
				g.Add(t.decodeRequest(signature, i))
				g.Add(t.decodeResponse(signature, i))
				g.Add(t.replyType(signature, i))
				g.Line()
			})
		f.Line()
	}

	return f
}

// Returns FieldTypeToProto.
func typeToProto(field *parser.FuncField) string {
	ptr, list := "", ""
	if field.IsPointer {
		ptr = "Ptr"
	}
	if field.IsArray {
		list = "List"
	}
	return util.ToUpperFirst(field.Type) + ptr + list + "ToProto"
}

// Returns ProtoToFieldType.
func protoToType(field *parser.FuncField) string {
	ptr, list := "", ""
	if field.IsPointer {
		ptr = "Ptr"
	}
	if field.IsArray {
		list = "List"
	}
	return "ProtoTo" + util.ToUpperFirst(field.Type) + ptr + list
}

func (GRPCConverterTemplate) Path() string {
	return "./transport/converter/protobuf/converters.go"
}

// Renders type conversion (if need) to default protobuf types.
//		req.Symbol
// or
//		int(resp.Count)
// or
//		structNamePositions
// based on field type
func golangTypeToProto(structName string, field *parser.FuncField) *Statement {
	if field.IsArray || field.IsPointer {
		return Id(structName + util.ToUpperFirst(field.Name))
	} else if isDefaultProtoField(field) {
		return Id(structName).Dot(util.ToUpperFirst(field.Name))
	} else if newType, ok := goToProtoTypesMap[field.Type]; ok {
		newField := &parser.FuncField{
			Type:      newType,
			Name:      field.Name,
			IsArray:   field.IsArray,
			Package:   field.Package,
			IsPointer: field.IsPointer,
		}
		return fieldType(newField).Call(Id(structName).Dot(util.ToUpperFirst(field.Name)))
	}
	return Id(structName + util.ToUpperFirst(field.Name))
}

func canConvertGolangToProto(field *parser.FuncField) bool {
	if field.IsArray || field.IsPointer {
		return false
	}
	if isDefaultProtoField(field) {
		return true
	}
	if _, ok := goToProtoTypesMap[field.Type]; ok {
		return true
	}
	return false
}

// Renders type conversion to default golang types.
// 		int(resp.Count)
// or
// 		structNamePositions
// based on field type
func protoTypeToGolang(structName string, field *parser.FuncField) *Statement {
	if field.IsArray || field.IsPointer {
		return Id(structName + util.ToUpperFirst(field.Name))
	} else if isDefaultGolangField(field) {
		return fieldType(field).Call(Id(structName).Dot(util.ToUpperFirst(field.Name)))
	}
	return Id(structName + util.ToUpperFirst(field.Name))
}

func canConvertProtoToGolang(field *parser.FuncField) bool {
	if field.IsArray || field.IsPointer {
		return false
	} else if isDefaultGolangField(field) {
		return true
	}
	return false
}

func isDefaultProtoField(field *parser.FuncField) bool {
	return util.IsInStringSlice(field.Type, defaultProtoTypes)
}

func isDefaultGolangField(field *parser.FuncField) bool {
	return util.IsInStringSlice(field.Type, defaultGolangTypes)
}

// Render custom type converting and error checking
//
//		structNamePositions, err := util.ProtoToIntList(structName.Positions)
//		if err != nil {
//			return nil, err
//		}
//
func (t GRPCConverterTemplate) convertCustomType(structName, converterName string, field *parser.FuncField) *Statement {
	return List(Id(structName+util.ToUpperFirst(field.Name)), Err()).
		Op(":=").
		Add(
			Qual(utilPackagePath(t.PackagePath), converterName).
				Call(Id(structName).
					Dot(util.ToUpperFirst(field.Name))),
		).
		Line().If(Err().Op("!=").Nil()).Block(
		Return().List(Nil(), Err()),
	)
}

// Renders function for encoding request, golang type converts to proto type.
//
//		func(_ context.Context, request interface{}) (interface{}, error) {
//			req := request.(*svc.CountRequest)
//			return &stringsvc.CountRequest{
//				Symbol: req.Symbol,
//				Text:   req.Text,
//			}, nil
//		}
//
func (t GRPCConverterTemplate) encodeRequest(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodParams := removeContextIfFirst(signature.Params)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("request").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodParams) > 0 {
				group.Id("req").Op(":=").Id("request").Assert(Op("*").Qual(t.PackagePath, requestStructName(signature)))
				for _, field := range methodParams {
					if !canConvertGolangToProto(field) {
						group.Add(t.convertCustomType("req", typeToProto(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(protobufPath(i), requestStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodParams {
					dict[structFieldName(field)] = Line().Add(golangTypeToProto("req", field))
				}
			})), Nil())
		},
	)
}

// Renders function for encoding response, golang type converts to proto type.
//
//		func(_ context.Context, response interface{}) (interface{}, error) {
//			resp := response.(*svc.CountResponse)
//			respPositions, err := util.IntListToProto(resp.Positions)
//			if err != nil {
//				return nil, err
//			}
//			return &stringsvc.CountResponse{
//				Count:     int64(resp.Count),
//				Positions: respPositions,
//			}, nil
//		}
//
func (t GRPCConverterTemplate) encodeResponse(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodResults := removeContextIfFirst(signature.Results)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("response").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodResults) > 0 {
				group.Id("resp").Op(":=").Id("response").Assert(Op("*").Qual(t.PackagePath, responseStructName(signature)))
				for _, field := range methodResults {
					if !canConvertGolangToProto(field) {
						group.Add(t.convertCustomType("resp", typeToProto(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(protobufPath(i), responseStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodResults {
					dict[structFieldName(field)] = Line().Add(golangTypeToProto("resp", field))
				}
			})), Nil())
		},
	)
}

// Renders function for decoding request, proto type converts to golang type.
//
//		func(_ context.Context, request interface{}) (interface{}, error) {
//			req := request.(*stringsvc.CountRequest)
//			return &svc.CountRequest{
//				Symbol: string(req.Symbol),
//				Text:   string(req.Text),
//			}, nil
//		}
//
func (t GRPCConverterTemplate) decodeRequest(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodParams := removeContextIfFirst(signature.Params)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("request").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodParams) > 0 {
				group.Id("req").Op(":=").Id("request").Assert(Op("*").Qual(protobufPath(i), requestStructName(signature)))
				for _, field := range methodParams {
					if !canConvertProtoToGolang(field) {
						group.Add(t.convertCustomType("req", protoToType(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(t.PackagePath, requestStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodParams {
					dict[structFieldName(field)] = Line().Add(protoTypeToGolang("req", field))
				}
			})), Nil())
		},
	)
}

// Renders function for decoding response, proto type converts to golang type.
//
//		func(_ context.Context, response interface{}) (interface{}, error) {
//			resp := response.(*stringsvc.CountResponse)
//			respPositions, err := util.ProtoToIntList(resp.Positions)
//			if err != nil {
//				return nil, err
//			}
//			return &svc.CountResponse{
//				Count:     int(resp.Count),
//				Positions: respPositions,
//			}, nil
//		}
//
func (t GRPCConverterTemplate) decodeResponse(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodResults := removeContextIfFirst(signature.Results)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("response").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodResults) > 0 {
				group.Id("resp").Op(":=").Id("response").Assert(Op("*").Qual(protobufPath(i), responseStructName(signature)))
				for _, field := range methodResults {
					if !canConvertProtoToGolang(field) {
						group.Add(t.convertCustomType("resp", protoToType(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(t.PackagePath, responseStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodResults {
					dict[structFieldName(field)] = Line().Add(protoTypeToGolang("resp", field))
				}
			})), Nil())
		},
	)
}

// Renders reply type argument
// 		stringsvc.CountResponse{}
func (t GRPCConverterTemplate) replyType(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	return Line().Qual(protobufPath(i), responseStructName(signature)).Values()
}
<commit_msg>refactor(template): remove unnecessary variables from typeToProto and protoToType methods<commit_after>package template

import (
	. "github.com/dave/jennifer/jen"
	"github.com/devimteam/microgen/parser"
	"github.com/devimteam/microgen/util"
)

var (
	defaultProtoTypes = []string{"string", "bool", "byte", "int64", "uint64", "float64", "int32", "uint32", "float32"}
	goToProtoTypesMap = map[string]string{
		"uint": "uint64",
		"int":  "int64",
	}
	defaultGolangTypes = []string{"string", "bool", "int", "uint", "byte", "int64", "uint64", "float64", "int32", "uint32", "float32"}
)

type GRPCConverterTemplate struct {
	PackagePath string
}

func utilPackagePath(path string) string {
	return path + "/util"
}

func (t GRPCClientTemplate) converterPackagePath() string {
	return t.PackagePath + "/transport/converter/protobuf"
}

// Renders converter file.
//
//		// This file was automatically generated by "microgen" utility.
//		// Please, do not edit.
//		package protobuf
//
//		import (
//			context "context"
//			grpc "github.com/devimteam/go-kit/transportlayer/grpc"
//			svc "github.com/devimteam/microgen/test/svc"
//			util "github.com/devimteam/microgen/test/svc/util"
//			stringsvc "gitlab.devim.team/protobuf/stringsvc"
//		)
//
//		var CountConverter = &grpc.EndpointConverter{
//			func(_ context.Context, request interface{}) (interface{}, error) {
//				req := request.(*svc.CountRequest)
//				return &stringsvc.CountRequest{
//					Symbol: req.Symbol,
//					Text:   req.Text,
//				}, nil
//			},
//			func(_ context.Context, response interface{}) (interface{}, error) {
//				resp := response.(*svc.CountResponse)
//				respPositions, err := util.IntListToProto(resp.Positions)
//				if err != nil {
//					return nil, err
//				}
//				return &stringsvc.CountResponse{
//					Count:     int64(resp.Count),
//					Positions: respPositions,
//				}, nil
//			},
//			func(_ context.Context, request interface{}) (interface{}, error) {
//				req := request.(*stringsvc.CountRequest)
//				return &svc.CountRequest{
//					Symbol: string(req.Symbol),
//					Text:   string(req.Text),
//				}, nil
//			},
//			func(_ context.Context, response interface{}) (interface{}, error) {
//				resp := response.(*stringsvc.CountResponse)
//				respPositions, err := util.ProtoToIntList(resp.Positions)
//				if err != nil {
//					return nil, err
//				}
//				return &svc.CountResponse{
//					Count:     int(resp.Count),
//					Positions: respPositions,
//				}, nil
//			},
//			stringsvc.CountResponse{},
//		}
//
func (t GRPCConverterTemplate) Render(i *parser.Interface) *File {
	f := NewFile("protobuf")

	for _, signature := range i.FuncSignatures {
		f.Var().Id(converterStructName(signature)).Op("=").Op("&").Qual(PackagePathTransportLayerGRPC, "EndpointConverter").
			ValuesFunc(func(g *Group) {
				g.Add(t.encodeRequest(signature, i))
				g.Add(t.encodeResponse(signature, i))
				g.Add(t.decodeRequest(signature, i))
				g.Add(t.decodeResponse(signature, i))
				g.Add(t.replyType(signature, i))
				g.Line()
			})
		f.Line()
	}

	return f
}

// Returns FieldTypeToProto.
func typeToProto(field *parser.FuncField) string {
	methodName := util.ToUpperFirst(field.Type)
	if field.IsPointer {
		methodName += "Ptr"
	}
	if field.IsArray {
		methodName += "List"
	}
	return methodName + "ToProto"
}

// Returns ProtoToFieldType.
func protoToType(field *parser.FuncField) string {
	methodName := "ProtoTo" + util.ToUpperFirst(field.Type)
	if field.IsPointer {
		methodName += "Ptr"
	}
	if field.IsArray {
		methodName += "List"
	}
	return methodName
}

func (GRPCConverterTemplate) Path() string {
	return "./transport/converter/protobuf/converters.go"
}

// Renders type conversion (if need) to default protobuf types.
//		req.Symbol
// or
//		int(resp.Count)
// or
//		structNamePositions
// based on field type
func golangTypeToProto(structName string, field *parser.FuncField) *Statement {
	if field.IsArray || field.IsPointer {
		return Id(structName + util.ToUpperFirst(field.Name))
	} else if isDefaultProtoField(field) {
		return Id(structName).Dot(util.ToUpperFirst(field.Name))
	} else if newType, ok := goToProtoTypesMap[field.Type]; ok {
		newField := &parser.FuncField{
			Type:      newType,
			Name:      field.Name,
			IsArray:   field.IsArray,
			Package:   field.Package,
			IsPointer: field.IsPointer,
		}
		return fieldType(newField).Call(Id(structName).Dot(util.ToUpperFirst(field.Name)))
	}
	return Id(structName + util.ToUpperFirst(field.Name))
}

func canConvertGolangToProto(field *parser.FuncField) bool {
	if field.IsArray || field.IsPointer {
		return false
	}
	if isDefaultProtoField(field) {
		return true
	}
	if _, ok := goToProtoTypesMap[field.Type]; ok {
		return true
	}
	return false
}

// Renders type conversion to default golang types.
// 		int(resp.Count)
// or
// 		structNamePositions
// based on field type
func protoTypeToGolang(structName string, field *parser.FuncField) *Statement {
	if field.IsArray || field.IsPointer {
		return Id(structName + util.ToUpperFirst(field.Name))
	} else if isDefaultGolangField(field) {
		return fieldType(field).Call(Id(structName).Dot(util.ToUpperFirst(field.Name)))
	}
	return Id(structName + util.ToUpperFirst(field.Name))
}

func canConvertProtoToGolang(field *parser.FuncField) bool {
	if field.IsArray || field.IsPointer {
		return false
	} else if isDefaultGolangField(field) {
		return true
	}
	return false
}

func isDefaultProtoField(field *parser.FuncField) bool {
	return util.IsInStringSlice(field.Type, defaultProtoTypes)
}

func isDefaultGolangField(field *parser.FuncField) bool {
	return util.IsInStringSlice(field.Type, defaultGolangTypes)
}

// Render custom type converting and error checking
//
//		structNamePositions, err := util.ProtoToIntList(structName.Positions)
//		if err != nil {
//			return nil, err
//		}
//
func (t GRPCConverterTemplate) convertCustomType(structName, converterName string, field *parser.FuncField) *Statement {
	return List(Id(structName+util.ToUpperFirst(field.Name)), Err()).
		Op(":=").
		Add(
			Qual(utilPackagePath(t.PackagePath), converterName).
				Call(Id(structName).
					Dot(util.ToUpperFirst(field.Name))),
		).
		Line().If(Err().Op("!=").Nil()).Block(
		Return().List(Nil(), Err()),
	)
}

// Renders function for encoding request, golang type converts to proto type.
//
//		func(_ context.Context, request interface{}) (interface{}, error) {
//			req := request.(*svc.CountRequest)
//			return &stringsvc.CountRequest{
//				Symbol: req.Symbol,
//				Text:   req.Text,
//			}, nil
//		}
//
func (t GRPCConverterTemplate) encodeRequest(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodParams := removeContextIfFirst(signature.Params)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("request").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodParams) > 0 {
				group.Id("req").Op(":=").Id("request").Assert(Op("*").Qual(t.PackagePath, requestStructName(signature)))
				for _, field := range methodParams {
					if !canConvertGolangToProto(field) {
						group.Add(t.convertCustomType("req", typeToProto(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(protobufPath(i), requestStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodParams {
					dict[structFieldName(field)] = Line().Add(golangTypeToProto("req", field))
				}
			})), Nil())
		},
	)
}

// Renders function for encoding response, golang type converts to proto type.
//
//		func(_ context.Context, response interface{}) (interface{}, error) {
//			resp := response.(*svc.CountResponse)
//			respPositions, err := util.IntListToProto(resp.Positions)
//			if err != nil {
//				return nil, err
//			}
//			return &stringsvc.CountResponse{
//				Count:     int64(resp.Count),
//				Positions: respPositions,
//			}, nil
//		}
//
func (t GRPCConverterTemplate) encodeResponse(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodResults := removeContextIfFirst(signature.Results)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("response").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodResults) > 0 {
				group.Id("resp").Op(":=").Id("response").Assert(Op("*").Qual(t.PackagePath, responseStructName(signature)))
				for _, field := range methodResults {
					if !canConvertGolangToProto(field) {
						group.Add(t.convertCustomType("resp", typeToProto(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(protobufPath(i), responseStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodResults {
					dict[structFieldName(field)] = Line().Add(golangTypeToProto("resp", field))
				}
			})), Nil())
		},
	)
}

// Renders function for decoding request, proto type converts to golang type.
//
//		func(_ context.Context, request interface{}) (interface{}, error) {
//			req := request.(*stringsvc.CountRequest)
//			return &svc.CountRequest{
//				Symbol: string(req.Symbol),
//				Text:   string(req.Text),
//			}, nil
//		}
//
func (t GRPCConverterTemplate) decodeRequest(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodParams := removeContextIfFirst(signature.Params)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("request").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodParams) > 0 {
				group.Id("req").Op(":=").Id("request").Assert(Op("*").Qual(protobufPath(i), requestStructName(signature)))
				for _, field := range methodParams {
					if !canConvertProtoToGolang(field) {
						group.Add(t.convertCustomType("req", protoToType(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(t.PackagePath, requestStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodParams {
					dict[structFieldName(field)] = Line().Add(protoTypeToGolang("req", field))
				}
			})), Nil())
		},
	)
}

// Renders function for decoding response, proto type converts to golang type.
//
//		func(_ context.Context, response interface{}) (interface{}, error) {
//			resp := response.(*stringsvc.CountResponse)
//			respPositions, err := util.ProtoToIntList(resp.Positions)
//			if err != nil {
//				return nil, err
//			}
//			return &svc.CountResponse{
//				Count:     int(resp.Count),
//				Positions: respPositions,
//			}, nil
//		}
//
func (t GRPCConverterTemplate) decodeResponse(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	methodResults := removeContextIfFirst(signature.Results)
	return Line().Func().Call(Op("_").Qual(PackagePathContext, "Context"), Id("response").Interface()).Params(Interface(), Error()).BlockFunc(
		func(group *Group) {
			if len(methodResults) > 0 {
				group.Id("resp").Op(":=").Id("response").Assert(Op("*").Qual(protobufPath(i), responseStructName(signature)))
				for _, field := range methodResults {
					if !canConvertProtoToGolang(field) {
						group.Add(t.convertCustomType("resp", protoToType(field), field))
					}
				}
			}
			group.Return().List(Op("&").Qual(t.PackagePath, responseStructName(signature)).Values(DictFunc(func(dict Dict) {
				for _, field := range methodResults {
					dict[structFieldName(field)] = Line().Add(protoTypeToGolang("resp", field))
				}
			})), Nil())
		},
	)
}

// Renders reply type argument
// 		stringsvc.CountResponse{}
func (t GRPCConverterTemplate) replyType(signature *parser.FuncSignature, i *parser.Interface) *Statement {
	return Line().Qual(protobufPath(i), responseStructName(signature)).Values()
}
<|endoftext|>


<commit_before>// Copyright 2016 AlexStocks(https://github.com/AlexStocks).
// All rights reserved.  Use of m source code is
// governed by a BSD-style license.
//
// 2016-09-11 19:30
// Package gxdriver provides a MySQL driver for Go's database/sql package
// code example: https://github.com/alexstocks/go-practice/blob/master/mysql/stmt.go
package gxdriver

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"log"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

import (
	"github.com/AlexStocks/goext/log"
	"github.com/AlexStocks/goext/strings"
	mysql "github.com/go-sql-driver/mysql"
)

const (
	GettyMySQLDriver = "getty_mysql_driver"
)

func init() {
	sql.Register(GettyMySQLDriver, &MySQLDriver{})
	dbConnMap = make(map[string]*sql.DB)
}

///////////////////////////////////////
// Mysql Driver
///////////////////////////////////////

// refs: https://zhuanlan.zhihu.com/p/24768377
// author: https://github.com/idada
// 重写driver，方便重用prepare statement.
// !: statement用到的connection和transaction begin的connection不是同一个connection
type mysqlStmt struct {
	driver.Stmt
	conn  *mySQLConn
	query string
	ref   int32
}

func (stmt *mysqlStmt) Close() error {
	stmt.conn.stmtMutex.Lock()
	defer stmt.conn.stmtMutex.Unlock()

	if atomic.AddInt32(&stmt.ref, -1) == 0 {
		delete(stmt.conn.stmtCache, stmt.query)
		return stmt.Stmt.Close()
	}
	return nil
}

type MySQLDriver struct {
}

func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
	gxlog.CInfo("GettyMSDriver:Open(%s)", dsn)
	var driver mysql.MySQLDriver
	conn, err := driver.Open(dsn)
	if err != nil {
		return nil, err
	}

	return &mySQLConn{Conn: conn, stmtCache: make(map[string]*mysqlStmt)}, nil
}

type mySQLConn struct {
	driver.Conn
	stmtMutex sync.RWMutex
	stmtCache map[string]*mysqlStmt
}

func (m *mySQLConn) Prepare(query string) (driver.Stmt, error) {
	gxlog.CInfo("GettyMSDriver:Prepare(%s)", query)
	m.stmtMutex.RLock()
	if stmt, exists := m.stmtCache[query]; exists {
		// must update reference counter in lock scope
		atomic.AddInt32(&stmt.ref, 1)
		m.stmtMutex.RUnlock()
		return stmt, nil
	}
	m.stmtMutex.RUnlock()

	m.stmtMutex.Lock()
	defer m.stmtMutex.Unlock()

	// double check
	if stmt, exists := m.stmtCache[query]; exists {
		atomic.AddInt32(&stmt.ref, 1)
		return stmt, nil
	}

	stmt, err := m.Conn.Prepare(query)
	if err != nil {
		return nil, err
	}

	stmt2 := &mysqlStmt{stmt, m, query, 1}
	m.stmtCache[query] = stmt2
	return stmt2, nil
}

func (m *mySQLConn) Begin() (driver.Tx, error) {
	gxlog.CInfo("GettyMSDriver:Begin")
	tx, err := m.Conn.Begin()
	if err != nil {
		return nil, err
	}
	return &mysqlTx{m, tx}, nil
}

type mysqlTx struct {
	*mySQLConn
	tx driver.Tx
}

func (tx *mysqlTx) Commit() (err error) {
	gxlog.CInfo("GettyMSDriver:Commit")
	return tx.tx.Commit()
}

func (tx *mysqlTx) Rollback() (err error) {
	return tx.tx.Rollback()
}

///////////////////////////////////////
// Mysql Instance
///////////////////////////////////////

var (
	dbConnMapLock sync.Mutex
	dbConnMap     map[string]*sql.DB
)

/*
feature list:
1 读写分离
2 长连接有效性检测和自动重连
3 mysql prepared statement 重用
! Note: 在执行sql操作之前请执行CheckSQl & CheckAcitve
*/
type MySQL struct {
	*sql.DB             //数据库操作
	tx          *sql.Tx //带事务的数据库操作
	txIndex     int     //事务记数器，只有当txIndex=0才会触发Begin动作，只有当txIndex=1才会触发Commit动作
	role        string  //操作类型，分Master或Slave
	schema      string  //数据库连接schema
	activetime  int64   //上次建立连接的时间点，需要一种机制来检测客户端与mysql服务端连接的有效性
	waittimeout int64   //mysql服务器空闲等待时长
	sync.Once
}

//创建一个默认的mysql操作实例
func Open(schema string) *MySQL {
	return newMySQLInstance(schema, "Master")
}

//创建一个默认的mysql查询实例
func OpenQuery(schema string) *MySQL {
	return newMySQLInstance(schema, "Slave")
}

func (m *MySQL) Close() {
	m.Do(func() {
		if m.tx != nil {
			m.tx.Commit()
		}
		m.tx = nil
		m.DB.Close()
		var key string = m.schema + m.role
		dbConnMapLock.Lock()
		db, ok := dbConnMap[key]
		dbConnMapLock.Unlock()
		if ok && db == m.DB {
			delete(dbConnMap, key)
		}
		m.DB = nil
	})
}

//@param string schema
//@param string role [Master, Slave]
func newMySQLInstance(schema string, role string) *MySQL {
	if !gxstrings.Contains([]string{"Master", "Slave"}, role) {
		panic("function common.NewSqlInstance's second argument must be 'Master' or 'Slave'.")
	}

	var key string = schema + role
	dbConnMapLock.Lock()
	conn, ok := dbConnMap[key]
	dbConnMapLock.Unlock()
	if !ok {
		//建立一个新连接到mysql
		connect(schema, role)
		dbConnMapLock.Lock()
		conn = dbConnMap[key]
		dbConnMapLock.Unlock()
	}
	return &MySQL{
		DB:         conn,
		schema:     schema,
		role:       role,
		activetime: time.Now().Unix(),
	}
}

//建立数据库连接
//@param string schema 连接DB方案
//@param string role 连接类型，是分Master和Slave类型
func connect(schema string, role string) {
	var key = schema + role
	//开始连接DB
	conn, err := sql.Open(GettyMySQLDriver, schema)
	if err != nil {
		log.Fatalln(err.Error())
	}

	//将DB连接放入一个全局变量中
	dbConnMapLock.Lock()
	dbConnMap[key] = conn
	dbConnMapLock.Unlock()
}

//根据db.Query的查询结果，组装成一个关联key的数据集，数据类型[]map[string]string
func (m *MySQL) FetchRows(rows *sql.Rows) ([]map[string]string, error) {
	result := make([]map[string]string, 0)
	columns, err := rows.Columns()
	if err != nil {
		//an error occurred
		return nil, err
	}

	rawBytes := make([]sql.RawBytes, len(columns))

	//rows.Scan wants '[]interface{}' as an argument, so we must copy
	//the references into such a slice
	scanArgs := make([]interface{}, len(columns))

	for i := range rawBytes {
		scanArgs[i] = &rawBytes[i]
	}

	for rows.Next() {
		err := rows.Scan(scanArgs...)
		if err != nil {
			return nil, err
		}
		var val string
		item := make(map[string]string)
		for i, col := range rawBytes {
			if col == nil {
				val = ""
			} else {
				val = string(col)
			}
			item[columns[i]] = val
		}
		result = append(result, item)
	}
	return result, nil
}

//检查MySQL实例的连接是否还在活跃时间范围内
func (m *MySQL) CheckActive() {
	var now int64 = time.Now().Unix()
	if m.tx != nil {
		//如果存在事务会话，则不再进行连接检查
		m.activetime = now
		return
	}

	//从MySQL的wait_timeout变量中定位waittimeout
	if m.waittimeout == 0 {
		rows, err := m.Query("SHOW VARIABLES LIKE 'wait_timeout'")
		if err != nil {
			log.Fatalln(err.Error())
		}
		defer rows.Close()

		result, err := m.FetchRows(rows)
		if err != nil {
			log.Fatalln(err.Error())
		}
		if result != nil && len(result) != 0 {
			timeout, err := strconv.Atoi(result[0]["Value"])
			if err != nil {
				log.Fatalln(err.Error())
			}
			m.waittimeout = int64(timeout)
		}
	}

	if now-m.activetime > m.waittimeout-2 {
		//此时认为数据库连接已经超时了，重新进行一次连接
		connect(m.schema, m.role)
		var key string = m.schema + m.role
		dbConnMapLock.Lock()
		m.DB = dbConnMap[key]
		dbConnMapLock.Unlock()
	}

	//设置当前时间为最新活跃点
	m.activetime = now
}

//保证修改、写入类的操作不在slave上执行
func (m *MySQL) CheckSQL(sql string) error {
	if m.role == "Slave" {
		sql = strings.TrimSpace(sql)
		exp := regexp.MustCompile(`^(?i:insert|update|delete|alter|truncate|drop)`)
		if exp.MatchString(sql) {
			return fmt.Errorf("insert|update|delete|alter|truncate|drop operation is not allowed on slave.")
		}
	}

	return nil
}

//开始一个事务，开始一个事务和提交、回滚事务必须一一对应
func (m *MySQL) Begin() error {
	if m.txIndex == 0 {
		var err error
		m.tx, err = m.DB.Begin()
		if err != nil {
			return err
		}
	}
	m.txIndex++

	return nil
}

//提交一个事务
func (m *MySQL) Commit() error {
	var err error
	if m.txIndex == 1 {
		err = m.tx.Commit()
		m.tx = nil
	}
	m.txIndex--

	return err
}

//事务回滚
func (m *MySQL) Rollback() error {
	err := m.tx.Rollback()
	m.txIndex = 0
	m.tx = nil

	return err
}
<commit_msg>fix mysql driver bug<commit_after>// Copyright 2016 AlexStocks(https://github.com/AlexStocks).
// All rights reserved.  Use of m source code is
// governed by a BSD-style license.
//
// 2016-09-11 19:30
// Package gxdriver provides a MySQL driver for Go's database/sql package
// code example: https://github.com/alexstocks/go-practice/blob/master/mysql/stmt.go
package gxdriver

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

import (
	// "github.com/AlexStocks/goext/log"
	"github.com/AlexStocks/goext/strings"
	mysql "github.com/go-sql-driver/mysql"
)

const (
	GettyMySQLDriver = "getty_mysql_driver"
)

func init() {
	sql.Register(GettyMySQLDriver, &MySQLDriver{})
	dbConnMap = make(map[string]*sql.DB)
}

///////////////////////////////////////
// Mysql Driver
///////////////////////////////////////

// refs: https://zhuanlan.zhihu.com/p/24768377
// author: https://github.com/idada
// 重写driver，方便重用prepare statement.
// !: statement用到的connection和transaction begin的connection不是同一个connection
type mysqlStmt struct {
	driver.Stmt
	conn  *mySQLConn
	query string
	ref   int32
}

func (stmt *mysqlStmt) Close() error {
	stmt.conn.stmtMutex.Lock()
	defer stmt.conn.stmtMutex.Unlock()

	if atomic.AddInt32(&stmt.ref, -1) == 0 {
		delete(stmt.conn.stmtCache, stmt.query)
		return stmt.Stmt.Close()
	}
	return nil
}

type MySQLDriver struct {
}

func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
	// gxlog.CInfo("GettyMSDriver:Open(%s)", dsn)
	var driver mysql.MySQLDriver
	conn, err := driver.Open(dsn)
	if err != nil {
		return nil, err
	}

	return &mySQLConn{Conn: conn, stmtCache: make(map[string]*mysqlStmt)}, nil
}

type mySQLConn struct {
	driver.Conn
	stmtMutex sync.RWMutex
	// 缓存prepared stmt链接，在程序启动的时候全都创建好
	stmtCache map[string]*mysqlStmt
}

func (m *mySQLConn) Prepare(query string) (driver.Stmt, error) {
	// gxlog.CInfo("GettyMSDriver:Prepare(%s)", query)
	m.stmtMutex.RLock()
	if stmt, exists := m.stmtCache[query]; exists {
		// must update reference counter in lock scope
		atomic.AddInt32(&stmt.ref, 1)
		m.stmtMutex.RUnlock()
		return stmt, nil
	}
	m.stmtMutex.RUnlock()

	m.stmtMutex.Lock()
	defer m.stmtMutex.Unlock()

	// double check
	if stmt, exists := m.stmtCache[query]; exists {
		atomic.AddInt32(&stmt.ref, 1)
		return stmt, nil
	}

	stmt, err := m.Conn.Prepare(query)
	if err != nil {
		return nil, err
	}

	stmt2 := &mysqlStmt{stmt, m, query, 1}
	m.stmtCache[query] = stmt2
	return stmt2, nil
}

func (m *mySQLConn) Begin() (driver.Tx, error) {
	// gxlog.CInfo("GettyMSDriver:Begin")
	tx, err := m.Conn.Begin()
	if err != nil {
		return nil, err
	}
	return &mysqlTx{m, tx}, nil
}

type mysqlTx struct {
	*mySQLConn
	tx driver.Tx
}

func (tx *mysqlTx) Commit() (err error) {
	// gxlog.CInfo("GettyMSDriver:Commit")
	return tx.tx.Commit()
}

func (tx *mysqlTx) Rollback() (err error) {
	return tx.tx.Rollback()
}

///////////////////////////////////////
// Mysql Instance
///////////////////////////////////////

var (
	dbConnMapLock sync.Mutex
	dbConnMap     map[string]*sql.DB
)

/*
feature list:
1 读写分离
2 长连接有效性检测和自动重连
3 mysql prepared statement 重用
! Note: 在执行sql操作之前请执行CheckSQl & CheckAcitve
*/
type MySQL struct {
	*sql.DB             //数据库操作
	tx          *sql.Tx //带事务的数据库操作
	txIndex     int     //事务记数器，只有当txIndex=0才会触发Begin动作，只有当txIndex=1才会触发Commit动作
	role        string  //操作类型，分Master或Slave
	schema      string  //数据库连接schema
	activetime  int64   //上次建立连接的时间点，需要一种机制来检测客户端与mysql服务端连接的有效性
	waittimeout int64   //mysql服务器空闲等待时长
	sync.Once
}

//创建一个默认的mysql操作实例
func Open(schema string) *MySQL {
	return newMySQLInstance(schema, "Master")
}

//创建一个默认的mysql查询实例
func OpenQuery(schema string) *MySQL {
	return newMySQLInstance(schema, "Slave")
}

func (m *MySQL) Close() {
	m.Do(func() {
		if m.tx != nil {
			m.tx.Commit()
		}
		m.tx = nil
		m.DB.Close()
		var key string = m.schema + m.role
		dbConnMapLock.Lock()
		db, ok := dbConnMap[key]
		if ok && db == m.DB {
			delete(dbConnMap, key)
		}
		dbConnMapLock.Unlock()
		m.DB = nil
	})
}

//@param string schema
//@param string role [Master, Slave]
func newMySQLInstance(schema string, role string) *MySQL {
	if !gxstrings.Contains([]string{"Master", "Slave"}, role) {
		panic("function common.NewSqlInstance's second argument must be 'Master' or 'Slave'.")
	}

	var key string = schema + role
	dbConnMapLock.Lock()
	conn, ok := dbConnMap[key]
	dbConnMapLock.Unlock()
	if !ok {
		//建立一个新连接到mysql
		connect(schema, role)
		dbConnMapLock.Lock()
		conn = dbConnMap[key]
		dbConnMapLock.Unlock()
	}
	return &MySQL{
		DB:         conn,
		schema:     schema,
		role:       role,
		activetime: time.Now().Unix(),
	}
}

//建立数据库连接
//@param string schema 连接DB方案
//@param string role 连接类型，是分Master和Slave类型
func connect(schema string, role string) {
	var key = schema + role
	//开始连接DB
	conn, err := sql.Open(GettyMySQLDriver, schema)
	if err != nil {
		panic(err)
	}

	//将DB连接放入一个全局变量中
	dbConnMapLock.Lock()
	dbConnMap[key] = conn
	dbConnMapLock.Unlock()
}

//根据db.Query的查询结果，组装成一个关联key的数据集，数据类型[]map[string]string
func (m *MySQL) FetchRows(rows *sql.Rows) ([]map[string]string, error) {
	result := make([]map[string]string, 0)
	columns, err := rows.Columns()
	if err != nil {
		//an error occurred
		return nil, err
	}

	rawBytes := make([]sql.RawBytes, len(columns))

	//rows.Scan wants '[]interface{}' as an argument, so we must copy
	//the references into such a slice
	scanArgs := make([]interface{}, len(columns))

	for i := range rawBytes {
		scanArgs[i] = &rawBytes[i]
	}

	for rows.Next() {
		err := rows.Scan(scanArgs...)
		if err != nil {
			return nil, err
		}
		var val string
		item := make(map[string]string)
		for i, col := range rawBytes {
			if col == nil {
				val = ""
			} else {
				val = string(col)
			}
			item[columns[i]] = val
		}
		result = append(result, item)
	}
	return result, nil
}

//检查MySQL实例的连接是否还在活跃时间范围内
func (m *MySQL) CheckActive() {
	var now int64 = time.Now().Unix()
	if m.tx != nil {
		//如果存在事务会话，则不再进行连接检查
		m.activetime = now
		return
	}

	//从MySQL的wait_timeout变量中定位waittimeout
	if m.waittimeout == 0 {
		rows, err := m.Query("SHOW VARIABLES LIKE 'wait_timeout'")
		if err != nil {
			panic(err)
		}
		defer rows.Close()

		result, err := m.FetchRows(rows)
		if err != nil {
			panic(err)
		}
		if result != nil && len(result) != 0 {
			timeout, err := strconv.Atoi(result[0]["Value"])
			if err != nil {
				panic(err)
			}
			m.waittimeout = int64(timeout)
		}
	}

	if now-m.activetime > m.waittimeout-2 {
		//此时认为数据库连接已经超时了，重新进行一次连接
		connect(m.schema, m.role)
		var key string = m.schema + m.role
		dbConnMapLock.Lock()
		m.DB = dbConnMap[key]
		dbConnMapLock.Unlock()
	}

	//设置当前时间为最新活跃点
	m.activetime = now
}

//保证修改、写入类的操作不在slave上执行
func (m *MySQL) CheckSQL(sql string) error {
	if m.role == "Slave" {
		sql = strings.TrimSpace(sql)
		exp := regexp.MustCompile(`^(?i:insert|update|delete|alter|truncate|drop)`)
		if exp.MatchString(sql) {
			return fmt.Errorf("insert|update|delete|alter|truncate|drop operation is not allowed on slave.")
		}
	}

	return nil
}

//开始一个事务，开始一个事务和提交、回滚事务必须一一对应
func (m *MySQL) Begin() error {
	if m.txIndex == 0 {
		var err error
		m.tx, err = m.DB.Begin()
		if err != nil {
			return err
		}
	}
	m.txIndex++

	return nil
}

//提交一个事务
func (m *MySQL) Commit() error {
	var err error
	if m.txIndex == 1 {
		err = m.tx.Commit()
		m.tx = nil
	}
	m.txIndex--

	return err
}

//事务回滚
func (m *MySQL) Rollback() error {
	err := m.tx.Rollback()
	m.txIndex = 0
	m.tx = nil

	return err
}
<|endoftext|>


<commit_before>package oskite

import (
	"fmt"
	"koding/virt"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	redigo "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

var (
	oskites   = make(map[string]*OskiteInfo)
	oskitesMu sync.Mutex
)

func (o *Oskite) loadBalancer(correlationName, username, deadService string) string {
	blog := func(v interface{}) {
		log.Info("oskite loadbalancer for [correlationName: '%s' user: '%s' deadService: '%s'] results in --> %v.", correlationName, username, deadService, v)
	}

	resultOskite := o.ServiceUniquename
	lowestOskite := lowestOskiteLoad()
	if lowestOskite != "" {
		if deadService == lowestOskite {
			resultOskite = o.ServiceUniquename
		} else {
			resultOskite = lowestOskite
		}
	}

	var vm *virt.VM
	if bson.IsObjectIdHex(correlationName) {
		mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.FindId(bson.ObjectIdHex(correlationName)).One(&vm)
		})
	}

	if vm == nil {
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Find(bson.M{"hostnameAlias": correlationName}).One(&vm)
		}); err != nil {
			blog(fmt.Sprintf("no hostnameAlias found, returning %s", resultOskite))
			return resultOskite // no vm was found, return this oskite
		}
	}

	if vm.PinnedToHost != "" {
		blog(fmt.Sprintf("returning pinnedHost '%s'", vm.PinnedToHost))
		return vm.PinnedToHost
	}

	if vm.HostKite == "" {
		// also set hoskite to prevent race condition between terminal and
		// oskite. Because if we set it now, the "kite.who" method of terminal
		// will not reply with an empty response.
		err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id, "hostKite": nil}, bson.M{"$set": bson.M{"hostKite": resultOskite}})
		})
		if err != nil {
			blog(fmt.Sprintf("hostkite is empty returning '%s' (update err: %s)", resultOskite, err))
		} else {
			blog(fmt.Sprintf("hostkite is empty returning '%s'", resultOskite))
		}
		return resultOskite
	}

	// maintenance and banned will be handled again in valideVM() function,
	// which will return a permission error.
	if vm.HostKite == "(maintenance)" || vm.HostKite == "(banned)" {
		blog(fmt.Sprintf("hostkite is %s returning '%s'", vm.HostKite, resultOskite))
		return resultOskite
	}

	// Set hostkite to nil if we detect a dead service. On the next call,
	// Oskite will point to an health service in validateVM function()
	// because it will detect that the hostkite is nil and change it to the
	// healthy service given by the client, which is the returned
	// k.ServiceUniqueName.
	if vm.HostKite == deadService {
		blog(fmt.Sprintf("dead service detected %s returning '%s'", vm.HostKite, o.ServiceUniquename))
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id}, bson.M{"$set": bson.M{"hostKite": nil}})
		}); err != nil {
			log.LogError(err, 0, vm.Id.Hex())
		}

		return resultOskite
	}

	blog(fmt.Sprintf("returning existing hostkite '%s'", vm.HostKite))
	return vm.HostKite
}

func (o *Oskite) setupRedis() {
	if o.RedisSession != nil {
		return
	}

	session, err := redis.NewRedisSession(conf.Redis)
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	o.RedisSession = session

	// oskite:production:sj:
	o.RedisPrefix = "oskite:" + conf.Environment + ":" + o.Region + ":"

	// kontainers-production-sj
	o.RedisKontainerSet = "kontainers-" + conf.Environment + "-" + o.Region

	// oskite:production:sj:kite-os-sj|kontainer3_sj_koding_com
	o.RedisKontainerKey = o.RedisPrefix + o.ServiceUniquename

	o.RedisSession.SetPrefix("oskite")
}

func (o *Oskite) redisBalancer() {
	o.setupRedis()

	log.Info("Connected to Redis with %s", o.RedisKontainerKey)
	_, err := redigo.Int(o.RedisSession.Do("SADD", o.RedisKontainerSet, o.RedisKontainerKey))
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	// update regularly our VMS info
	go func() {
		expireDuration := time.Second * 5
		for _ = range time.Tick(2 * time.Second) {
			oskiteInfo := o.GetOskiteInfo()

			if _, err := o.RedisSession.Do("HMSET", redigo.Args{o.RedisKontainerKey}.AddFlat(oskiteInfo)...); err != nil {
				log.Error("redis HMSET err: %v", err.Error())
			}

			reply, err := redigo.Int(o.RedisSession.Do("EXPIRE", o.RedisKontainerKey, expireDuration.Seconds()))
			if err != nil {
				log.Error("redis SET Expire %v. reply: %v err: %v", o.RedisKontainerKey, reply, err.Error())
			}
		}
	}()

	// get oskite statuses from others every 2 seconds
	for _ = range time.Tick(2 * time.Second) {
		kontainers, err := redigo.Strings(o.RedisSession.Do("SMEMBERS", o.RedisKontainerSet))
		if err != nil {
			log.Error("redis SMEMBER kontainers. err: %v", err.Error())
		}

		for _, kontainerHostname := range kontainers {
			// convert to o.ServiceUniquename format
			remoteOskite := strings.TrimPrefix(kontainerHostname, o.RedisPrefix)

			values, err := redigo.Values(o.RedisSession.Do("HGETALL", kontainerHostname))
			if err != nil {
				log.Error("redis HTGETALL %s. err: %v", kontainerHostname, err.Error())

				// kontainer might be dead, key gets than expired, continue with the next one

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			// cleanup members from the map if the key expires. Usually that
			// might due a host who added itself to the kontainer set but then
			// just died or added a wrong set name. If the kontainer comes up
			// again we do add it below again.
			if len(values) == 0 {
				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			oskiteInfo := new(OskiteInfo)
			if err := redigo.ScanStruct(values, oskiteInfo); err != nil {
				log.Error("redis ScanStruct err: %v", err.Error())
			}

			oskitesMu.Lock()
			oskites[remoteOskite] = oskiteInfo
			oskitesMu.Unlock()
		}
	}
}

func lowestOskiteLoad() (serviceUniquename string) {
	oskitesMu.Lock()
	defer oskitesMu.Unlock()

	oskitesSlice := make([]*OskiteInfo, 0, len(oskites))

	for s, v := range oskites {
		v.ServiceUniquename = s
		oskitesSlice = append(oskitesSlice, v)
	}

	sort.Sort(ByVM(oskitesSlice))

	middle := len(oskitesSlice) / 2
	if middle == 0 {
		return ""
	}

	// return randomly one of the lowest
	l := oskitesSlice[rand.Intn(middle)]

	// also pick up the highest to log information
	h := oskitesSlice[len(oskitesSlice)-1]

	log.Info("oskite picked up as lowest load %s with %d VMs (highest was: %d / %s)",
		l.ServiceUniquename, l.ActiveVMs, h.ActiveVMs, h.ServiceUniquename)

	return l.ServiceUniquename
}
<commit_msg>Ensure we have write safety enabled to avoid inconsistent reads elsewhere in the kite<commit_after>package oskite

import (
	"fmt"
	"koding/virt"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	redigo "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

var (
	oskites   = make(map[string]*OskiteInfo)
	oskitesMu sync.Mutex
)

func (o *Oskite) loadBalancer(correlationName, username, deadService string) string {
	blog := func(v interface{}) {
		log.Info("oskite loadbalancer for [correlationName: '%s' user: '%s' deadService: '%s'] results in --> %v.", correlationName, username, deadService, v)
	}
	// Ensure we are using a mongo master so that we can avoid db induced races
	mongodbConn.Session.SetSafe(&mgo.Safe{W: 3, WTimeout: 2, FSync: true})

	resultOskite := o.ServiceUniquename
	lowestOskite := lowestOskiteLoad()
	if lowestOskite != "" {
		if deadService == lowestOskite {
			resultOskite = o.ServiceUniquename
		} else {
			resultOskite = lowestOskite
		}
	}

	var vm *virt.VM
	if bson.IsObjectIdHex(correlationName) {
		mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.FindId(bson.ObjectIdHex(correlationName)).One(&vm)
		})
	}

	if vm == nil {
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Find(bson.M{"hostnameAlias": correlationName}).One(&vm)
		}); err != nil {
			blog(fmt.Sprintf("no hostnameAlias found, returning %s", resultOskite))
			return resultOskite // no vm was found, return this oskite
		}
	}

	if vm.PinnedToHost != "" {
		blog(fmt.Sprintf("returning pinnedHost '%s'", vm.PinnedToHost))
		return vm.PinnedToHost
	}

	if vm.HostKite == "" {
		// also set hoskite to prevent race condition between terminal and
		// oskite. Because if we set it now, the "kite.who" method of terminal
		// will not reply with an empty response.
		err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id, "hostKite": nil}, bson.M{"$set": bson.M{"hostKite": resultOskite}})
		})
		if err != nil {
			blog(fmt.Sprintf("hostkite is empty returning '%s' (update err: %s)", resultOskite, err))
		} else {
			blog(fmt.Sprintf("hostkite is empty returning '%s'", resultOskite))
		}
		return resultOskite
	}

	// maintenance and banned will be handled again in valideVM() function,
	// which will return a permission error.
	if vm.HostKite == "(maintenance)" || vm.HostKite == "(banned)" {
		blog(fmt.Sprintf("hostkite is %s returning '%s'", vm.HostKite, resultOskite))
		return resultOskite
	}

	// Set hostkite to nil if we detect a dead service. On the next call,
	// Oskite will point to an health service in validateVM function()
	// because it will detect that the hostkite is nil and change it to the
	// healthy service given by the client, which is the returned
	// k.ServiceUniqueName.
	if vm.HostKite == deadService {
		blog(fmt.Sprintf("dead service detected %s returning '%s'", vm.HostKite, o.ServiceUniquename))
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id}, bson.M{"$set": bson.M{"hostKite": nil}})
		}); err != nil {
			log.LogError(err, 0, vm.Id.Hex())
		}

		return resultOskite
	}

	blog(fmt.Sprintf("returning existing hostkite '%s'", vm.HostKite))
	return vm.HostKite
}

func (o *Oskite) setupRedis() {
	if o.RedisSession != nil {
		return
	}

	session, err := redis.NewRedisSession(conf.Redis)
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	o.RedisSession = session

	// oskite:production:sj:
	o.RedisPrefix = "oskite:" + conf.Environment + ":" + o.Region + ":"

	// kontainers-production-sj
	o.RedisKontainerSet = "kontainers-" + conf.Environment + "-" + o.Region

	// oskite:production:sj:kite-os-sj|kontainer3_sj_koding_com
	o.RedisKontainerKey = o.RedisPrefix + o.ServiceUniquename

	o.RedisSession.SetPrefix("oskite")
}

func (o *Oskite) redisBalancer() {
	o.setupRedis()

	log.Info("Connected to Redis with %s", o.RedisKontainerKey)
	_, err := redigo.Int(o.RedisSession.Do("SADD", o.RedisKontainerSet, o.RedisKontainerKey))
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	// update regularly our VMS info
	go func() {
		expireDuration := time.Second * 5
		for _ = range time.Tick(2 * time.Second) {
			oskiteInfo := o.GetOskiteInfo()

			if _, err := o.RedisSession.Do("HMSET", redigo.Args{o.RedisKontainerKey}.AddFlat(oskiteInfo)...); err != nil {
				log.Error("redis HMSET err: %v", err.Error())
			}

			reply, err := redigo.Int(o.RedisSession.Do("EXPIRE", o.RedisKontainerKey, expireDuration.Seconds()))
			if err != nil {
				log.Error("redis SET Expire %v. reply: %v err: %v", o.RedisKontainerKey, reply, err.Error())
			}
		}
	}()

	// get oskite statuses from others every 2 seconds
	for _ = range time.Tick(2 * time.Second) {
		kontainers, err := redigo.Strings(o.RedisSession.Do("SMEMBERS", o.RedisKontainerSet))
		if err != nil {
			log.Error("redis SMEMBER kontainers. err: %v", err.Error())
		}

		for _, kontainerHostname := range kontainers {
			// convert to o.ServiceUniquename format
			remoteOskite := strings.TrimPrefix(kontainerHostname, o.RedisPrefix)

			values, err := redigo.Values(o.RedisSession.Do("HGETALL", kontainerHostname))
			if err != nil {
				log.Error("redis HTGETALL %s. err: %v", kontainerHostname, err.Error())

				// kontainer might be dead, key gets than expired, continue with the next one

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			// cleanup members from the map if the key expires. Usually that
			// might due a host who added itself to the kontainer set but then
			// just died or added a wrong set name. If the kontainer comes up
			// again we do add it below again.
			if len(values) == 0 {
				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			oskiteInfo := new(OskiteInfo)
			if err := redigo.ScanStruct(values, oskiteInfo); err != nil {
				log.Error("redis ScanStruct err: %v", err.Error())
			}

			oskitesMu.Lock()
			oskites[remoteOskite] = oskiteInfo
			oskitesMu.Unlock()
		}
	}
}

func lowestOskiteLoad() (serviceUniquename string) {
	oskitesMu.Lock()
	defer oskitesMu.Unlock()

	oskitesSlice := make([]*OskiteInfo, 0, len(oskites))

	for s, v := range oskites {
		v.ServiceUniquename = s
		oskitesSlice = append(oskitesSlice, v)
	}

	sort.Sort(ByVM(oskitesSlice))

	middle := len(oskitesSlice) / 2
	if middle == 0 {
		return ""
	}

	// return randomly one of the lowest
	l := oskitesSlice[rand.Intn(middle)]

	// also pick up the highest to log information
	h := oskitesSlice[len(oskitesSlice)-1]

	log.Info("oskite picked up as lowest load %s with %d VMs (highest was: %d / %s)",
		l.ServiceUniquename, l.ActiveVMs, h.ActiveVMs, h.ServiceUniquename)

	return l.ServiceUniquename
}
<|endoftext|>


<commit_before>package oskite

import (
	"fmt"
	"koding/virt"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	redigo "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

var (
	oskites   = make(map[string]*OskiteInfo)
	oskitesMu sync.Mutex
)

func (o *Oskite) loadBalancer(correlationName, username, deadService string) string {
	blog := func(v interface{}) {
		log.Info("oskite loadbalancer for [correlationName: '%s' user: '%s' deadService: '%s'] results in --> %v.", correlationName, username, deadService, v)
	}

	resultOskite := o.ServiceUniquename
	lowestOskite := lowestOskiteLoad()
	if lowestOskite != "" {
		if deadService == lowestOskite {
			resultOskite = o.ServiceUniquename
		} else {
			resultOskite = lowestOskite
		}
	}

	var vm *virt.VM
	if bson.IsObjectIdHex(correlationName) {
		mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.FindId(bson.ObjectIdHex(correlationName)).One(&vm)
		})
	}

	if vm == nil {
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Find(bson.M{"hostnameAlias": correlationName}).One(&vm)
		}); err != nil {
			blog(fmt.Sprintf("no hostnameAlias found, returning %s", resultOskite))
			return resultOskite // no vm was found, return this oskite
		}
	}

	if vm.PinnedToHost != "" {
		blog(fmt.Sprintf("returning pinnedHost '%s'", vm.PinnedToHost))
		return vm.PinnedToHost
	}

	if vm.HostKite == "" {
		// also set hoskite to prevent race condition between terminal and
		// oskite. Because if we set it now, the "kite.who" method of terminal
		// will not reply with an empty response.
		err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id, "hostKite": nil}, bson.M{"$set": bson.M{"hostKite": resultOskite}})
		})

		blog(fmt.Sprintf("hostkite is empty returning '%s. (update err: %s)", resultOskite, err))
		return resultOskite
	}

	// maintenance and banned will be handled again in valideVM() function,
	// which will return a permission error.
	if vm.HostKite == "(maintenance)" || vm.HostKite == "(banned)" {
		blog(fmt.Sprintf("hostkite is %s returning '%s'", vm.HostKite, resultOskite))
		return resultOskite
	}

	// Set hostkite to nil if we detect a dead service. On the next call,
	// Oskite will point to an health service in validateVM function()
	// because it will detect that the hostkite is nil and change it to the
	// healthy service given by the client, which is the returned
	// k.ServiceUniqueName.
	if vm.HostKite == deadService {
		blog(fmt.Sprintf("dead service detected %s returning '%s'", vm.HostKite, o.ServiceUniquename))
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id}, bson.M{"$set": bson.M{"hostKite": nil}})
		}); err != nil {
			log.LogError(err, 0, vm.Id.Hex())
		}

		return resultOskite
	}

	blog(fmt.Sprintf("returning existing hostkite '%s'", vm.HostKite))
	return vm.HostKite
}

func (o *Oskite) setupRedis() {
	if o.RedisSession != nil {
		return
	}

	session, err := redis.NewRedisSession(conf.Redis)
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	o.RedisSession = session

	// oskite:production:sj:
	o.RedisPrefix = "oskite:" + conf.Environment + ":" + o.Region + ":"

	// kontainers-production-sj
	o.RedisKontainerSet = "kontainers-" + conf.Environment + "-" + o.Region

	// oskite:production:sj:kite-os-sj|kontainer3_sj_koding_com
	o.RedisKontainerKey = o.RedisPrefix + o.ServiceUniquename

	o.RedisSession.SetPrefix("oskite")
}

func (o *Oskite) redisBalancer() {
	o.setupRedis()

	log.Info("Connected to Redis with %s", o.RedisKontainerKey)
	_, err := redigo.Int(o.RedisSession.Do("SADD", o.RedisKontainerSet, o.RedisKontainerKey))
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	// update regularly our VMS info
	go func() {
		expireDuration := time.Second * 5
		for _ = range time.Tick(2 * time.Second) {
			oskiteInfo := o.GetOskiteInfo()

			if _, err := o.RedisSession.Do("HMSET", redigo.Args{o.RedisKontainerKey}.AddFlat(oskiteInfo)...); err != nil {
				log.Error("redis HMSET err: %v", err.Error())
			}

			reply, err := redigo.Int(o.RedisSession.Do("EXPIRE", o.RedisKontainerKey, expireDuration.Seconds()))
			if err != nil {
				log.Error("redis SET Expire %v. reply: %v err: %v", o.RedisKontainerKey, reply, err.Error())
			}
		}
	}()

	// get oskite statuses from others every 2 seconds
	for _ = range time.Tick(2 * time.Second) {
		kontainers, err := redigo.Strings(o.RedisSession.Do("SMEMBERS", o.RedisKontainerSet))
		if err != nil {
			log.Error("redis SMEMBER kontainers. err: %v", err.Error())
		}

		for _, kontainerHostname := range kontainers {
			// convert to o.ServiceUniquename format
			remoteOskite := strings.TrimPrefix(kontainerHostname, o.RedisPrefix)

			values, err := redigo.Values(o.RedisSession.Do("HGETALL", kontainerHostname))
			if err != nil {
				log.Error("redis HTGETALL %s. err: %v", kontainerHostname, err.Error())

				// kontainer might be dead, key gets than expired, continue with the next one

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			// cleanup members from the map if the key expires. Usually that
			// might due a host who added itself to the kontainer set but then
			// just died or added a wrong set name. If the kontainer comes up
			// again we do add it below again.
			if len(values) == 0 {
				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			oskiteInfo := new(OskiteInfo)
			if err := redigo.ScanStruct(values, oskiteInfo); err != nil {
				log.Error("redis ScanStruct err: %v", err.Error())
			}

			oskitesMu.Lock()
			oskites[remoteOskite] = oskiteInfo
			oskitesMu.Unlock()
		}
	}
}

func lowestOskiteLoad() (serviceUniquename string) {
	oskitesMu.Lock()
	defer oskitesMu.Unlock()

	oskitesSlice := make([]*OskiteInfo, 0, len(oskites))

	for s, v := range oskites {
		v.ServiceUniquename = s
		oskitesSlice = append(oskitesSlice, v)
	}

	sort.Sort(ByVM(oskitesSlice))

	middle := len(oskitesSlice) / 2
	if middle == 0 {
		return ""
	}

	// return randomly one of the lowest
	l := oskitesSlice[rand.Intn(middle)]

	// also pick up the highest to log information
	h := oskitesSlice[len(oskitesSlice)-1]

	log.Info("oskite picked up as lowest load %s with %d VMs (highest was: %d / %s)",
		l.ServiceUniquename, l.ActiveVMs, h.ActiveVMs, h.ServiceUniquename)

	return l.ServiceUniquename
}
<commit_msg>Only log the error if there is an error<commit_after>package oskite

import (
	"fmt"
	"koding/virt"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	redigo "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

var (
	oskites   = make(map[string]*OskiteInfo)
	oskitesMu sync.Mutex
)

func (o *Oskite) loadBalancer(correlationName, username, deadService string) string {
	blog := func(v interface{}) {
		log.Info("oskite loadbalancer for [correlationName: '%s' user: '%s' deadService: '%s'] results in --> %v.", correlationName, username, deadService, v)
	}

	resultOskite := o.ServiceUniquename
	lowestOskite := lowestOskiteLoad()
	if lowestOskite != "" {
		if deadService == lowestOskite {
			resultOskite = o.ServiceUniquename
		} else {
			resultOskite = lowestOskite
		}
	}

	var vm *virt.VM
	if bson.IsObjectIdHex(correlationName) {
		mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.FindId(bson.ObjectIdHex(correlationName)).One(&vm)
		})
	}

	if vm == nil {
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Find(bson.M{"hostnameAlias": correlationName}).One(&vm)
		}); err != nil {
			blog(fmt.Sprintf("no hostnameAlias found, returning %s", resultOskite))
			return resultOskite // no vm was found, return this oskite
		}
	}

	if vm.PinnedToHost != "" {
		blog(fmt.Sprintf("returning pinnedHost '%s'", vm.PinnedToHost))
		return vm.PinnedToHost
	}

	if vm.HostKite == "" {
		// also set hoskite to prevent race condition between terminal and
		// oskite. Because if we set it now, the "kite.who" method of terminal
		// will not reply with an empty response.
		err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id, "hostKite": nil}, bson.M{"$set": bson.M{"hostKite": resultOskite}})
		})
		if err != nil {
			blog(fmt.Sprintf("hostkite is empty returning '%s' (update err: %s)", resultOskite, err))
		} else {
			blog(fmt.Sprintf("hostkite is empty returning '%s'", resultOskite))
		}
		return resultOskite
	}

	// maintenance and banned will be handled again in valideVM() function,
	// which will return a permission error.
	if vm.HostKite == "(maintenance)" || vm.HostKite == "(banned)" {
		blog(fmt.Sprintf("hostkite is %s returning '%s'", vm.HostKite, resultOskite))
		return resultOskite
	}

	// Set hostkite to nil if we detect a dead service. On the next call,
	// Oskite will point to an health service in validateVM function()
	// because it will detect that the hostkite is nil and change it to the
	// healthy service given by the client, which is the returned
	// k.ServiceUniqueName.
	if vm.HostKite == deadService {
		blog(fmt.Sprintf("dead service detected %s returning '%s'", vm.HostKite, o.ServiceUniquename))
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id}, bson.M{"$set": bson.M{"hostKite": nil}})
		}); err != nil {
			log.LogError(err, 0, vm.Id.Hex())
		}

		return resultOskite
	}

	blog(fmt.Sprintf("returning existing hostkite '%s'", vm.HostKite))
	return vm.HostKite
}

func (o *Oskite) setupRedis() {
	if o.RedisSession != nil {
		return
	}

	session, err := redis.NewRedisSession(conf.Redis)
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	o.RedisSession = session

	// oskite:production:sj:
	o.RedisPrefix = "oskite:" + conf.Environment + ":" + o.Region + ":"

	// kontainers-production-sj
	o.RedisKontainerSet = "kontainers-" + conf.Environment + "-" + o.Region

	// oskite:production:sj:kite-os-sj|kontainer3_sj_koding_com
	o.RedisKontainerKey = o.RedisPrefix + o.ServiceUniquename

	o.RedisSession.SetPrefix("oskite")
}

func (o *Oskite) redisBalancer() {
	o.setupRedis()

	log.Info("Connected to Redis with %s", o.RedisKontainerKey)
	_, err := redigo.Int(o.RedisSession.Do("SADD", o.RedisKontainerSet, o.RedisKontainerKey))
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	// update regularly our VMS info
	go func() {
		expireDuration := time.Second * 5
		for _ = range time.Tick(2 * time.Second) {
			oskiteInfo := o.GetOskiteInfo()

			if _, err := o.RedisSession.Do("HMSET", redigo.Args{o.RedisKontainerKey}.AddFlat(oskiteInfo)...); err != nil {
				log.Error("redis HMSET err: %v", err.Error())
			}

			reply, err := redigo.Int(o.RedisSession.Do("EXPIRE", o.RedisKontainerKey, expireDuration.Seconds()))
			if err != nil {
				log.Error("redis SET Expire %v. reply: %v err: %v", o.RedisKontainerKey, reply, err.Error())
			}
		}
	}()

	// get oskite statuses from others every 2 seconds
	for _ = range time.Tick(2 * time.Second) {
		kontainers, err := redigo.Strings(o.RedisSession.Do("SMEMBERS", o.RedisKontainerSet))
		if err != nil {
			log.Error("redis SMEMBER kontainers. err: %v", err.Error())
		}

		for _, kontainerHostname := range kontainers {
			// convert to o.ServiceUniquename format
			remoteOskite := strings.TrimPrefix(kontainerHostname, o.RedisPrefix)

			values, err := redigo.Values(o.RedisSession.Do("HGETALL", kontainerHostname))
			if err != nil {
				log.Error("redis HTGETALL %s. err: %v", kontainerHostname, err.Error())

				// kontainer might be dead, key gets than expired, continue with the next one

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			// cleanup members from the map if the key expires. Usually that
			// might due a host who added itself to the kontainer set but then
			// just died or added a wrong set name. If the kontainer comes up
			// again we do add it below again.
			if len(values) == 0 {
				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			oskiteInfo := new(OskiteInfo)
			if err := redigo.ScanStruct(values, oskiteInfo); err != nil {
				log.Error("redis ScanStruct err: %v", err.Error())
			}

			oskitesMu.Lock()
			oskites[remoteOskite] = oskiteInfo
			oskitesMu.Unlock()
		}
	}
}

func lowestOskiteLoad() (serviceUniquename string) {
	oskitesMu.Lock()
	defer oskitesMu.Unlock()

	oskitesSlice := make([]*OskiteInfo, 0, len(oskites))

	for s, v := range oskites {
		v.ServiceUniquename = s
		oskitesSlice = append(oskitesSlice, v)
	}

	sort.Sort(ByVM(oskitesSlice))

	middle := len(oskitesSlice) / 2
	if middle == 0 {
		return ""
	}

	// return randomly one of the lowest
	l := oskitesSlice[rand.Intn(middle)]

	// also pick up the highest to log information
	h := oskitesSlice[len(oskitesSlice)-1]

	log.Info("oskite picked up as lowest load %s with %d VMs (highest was: %d / %s)",
		l.ServiceUniquename, l.ActiveVMs, h.ActiveVMs, h.ServiceUniquename)

	return l.ServiceUniquename
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	Aws struct {
		Key    string
		Secret string
	}
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     RuntimeOptions
	}
	Mongo string
	Mq    struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	ProxyKite         struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	Slack struct {
		Token   string
		Channel string
	}
	LogLevel             map[string]string
	Redis                string
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
	SocialApi struct {
		ProxyUrl     string
		CustomDomain struct {
			Public string
			Local  string
		}
	}
	Vmwatcher struct {
		Port           string
		AwsKey         string
		AwsSecret      string
		KloudSecretKey string
		KloudAddr      string
	}
	Segment        string
	GatherIngestor struct {
		Port int
	}
	Mailgun struct {
		Domain     string
		PrivateKey string
		PublicKey  string
	}
}

type RuntimeOptions struct {
	Kites struct {
		DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
		Kontrol                   struct {
			Username string `json:"username"`
		} `json:"kontrol"`
		Os struct {
			Version string `json:"version"`
		} `json:"os"`
		Terminal struct {
			Version string `json:"version"`
		} `json:"terminal"`
		Klient struct {
			Version string `json:"version"`
		} `json:"klient"`
		Kloud struct {
			Version string `json:"version"`
		} `json:"kloud"`
	} `json:"kites"`
	Algolia struct {
		AppId       string `json:"appId"`
		ApiKey      string `json:"apiKey"`
		IndexSuffix string `json:"indexSuffix"`
	} `json:"algolia"`
	SuppressLogs    bool   `json:"suppressLogs"`
	AuthExchange    string `json:"authExchange"`
	Environment     string `json:"environment"`
	Version         string `json:"version"`
	ResourceName    string `json:"resourceName"`
	UserSitesDomain string `json:"userSitesDomain"`
	LogResourceName string `json:"logResourceName"`
	SocialApiUri    string `json:"socialApiUri"`
	ApiUri          string `json:"apiUri"`
	MainUri         string `json:"mainUri"`
	SourceMapsUri   string `json:"sourceMapsUri"`
	Broker          struct {
		Uri string `json:"uri"`
	} `json:"broker"`
	AppsUri            string `json:"appsUri"`
	UploadsUri         string `json:"uploadsUri"`
	UploadsUriForGroup string `json:"uploadsUriForGroup"`
	FileFetchTimeout   int    `json:"fileFetchTimeout"`
	UserIdleMs         int    `json:"userIdleMs"`
	Embedly            struct {
		ApiKey string `json:"apiKey"`
	} `json:"embedly"`
	Github struct {
		ClientId string `json:"clientId"`
	} `json:"github"`
	Newkontrol struct {
		Url string `json:"url"`
	} `json:"newkontrol"`
	SessionCookie struct {
		MaxAge int  `json:"maxAge"`
		Secure bool `json:"secure"`
	} `json:"sessionCookie"`
	Stripe struct {
		Token string `json:"token"`
	} `json:"stripe"`
	ExternalProfiles struct {
		Google struct {
			Nicename string `json:"nicename"`
		} `json:"google"`
		Linkedin struct {
			Nicename string `json:"nicename"`
		} `json:"linkedin"`
		Twitter struct {
			Nicename string `json:"nicename"`
		} `json:"twitter"`
		Odesk struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"odesk"`
		Facebook struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"facebook"`
		Github struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"github"`
	} `json:"externalProfiles"`
	EntryPoint struct {
		Slug string `json:"slug"`
		Type string `json:"type"`
	} `json:"entryPoint"`
	Roles       []string      `json:"roles"`
	Permissions []interface{} `json:"permissions"`
	SiftScience string        `json:"siftScience"`
	Paypal      struct {
		FormUrl string `json:"formUrl"`
	} `json:"paypal"`
	Pubnub struct {
		SubscribeKey string `json:"subscribekey"`
		Enabled      bool   `json:"enabled"`
		SSL          bool   `json:"ssl"`
	} `json:"pubnub"`
	Collaboration struct {
		Timeout int `json:"timeout"`
	} `json:"collaboration"`
	PaymentBlockDuration float64 `json:"paymentBlockDuration"`
	DisabledFeatures     struct {
		Moderation bool `json:"moderation"`
		Teams      bool `json:"teams"`
		BotChannel bool `json:"botchannel"`
	} `json:"disabledFeatures"`
	ContentRotatorUrl string `json:"contentRotatorUrl"`
	Integration       struct {
		Url string `json:"url"`
	} `json:"integration"`
	WebhookMiddleware struct {
		Url string `json:"url"`
	} `json:"webhookMiddleware"`
	Google struct {
		ApiKey string `json:"apiKey"`
	} `json:"google"`
	Recaptcha struct {
		Key     string `json:"key"`
		Enabled bool   `json:"enabled"`
	} `json:"recaptcha"`
	Domains struct {
		Base string `json:"base"`
		Mail string `json:"mail"`
		Main string `json:"main"`
		Port string `json:"port"`
	} `json:"domains"`
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<commit_msg>go/config: remove odesk<commit_after>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	Aws struct {
		Key    string
		Secret string
	}
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     RuntimeOptions
	}
	Mongo string
	Mq    struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	ProxyKite         struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	Slack struct {
		Token   string
		Channel string
	}
	LogLevel             map[string]string
	Redis                string
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
	SocialApi struct {
		ProxyUrl     string
		CustomDomain struct {
			Public string
			Local  string
		}
	}
	Vmwatcher struct {
		Port           string
		AwsKey         string
		AwsSecret      string
		KloudSecretKey string
		KloudAddr      string
	}
	Segment        string
	GatherIngestor struct {
		Port int
	}
	Mailgun struct {
		Domain     string
		PrivateKey string
		PublicKey  string
	}
}

type RuntimeOptions struct {
	Kites struct {
		DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
		Kontrol                   struct {
			Username string `json:"username"`
		} `json:"kontrol"`
		Os struct {
			Version string `json:"version"`
		} `json:"os"`
		Terminal struct {
			Version string `json:"version"`
		} `json:"terminal"`
		Klient struct {
			Version string `json:"version"`
		} `json:"klient"`
		Kloud struct {
			Version string `json:"version"`
		} `json:"kloud"`
	} `json:"kites"`
	Algolia struct {
		AppId       string `json:"appId"`
		ApiKey      string `json:"apiKey"`
		IndexSuffix string `json:"indexSuffix"`
	} `json:"algolia"`
	SuppressLogs    bool   `json:"suppressLogs"`
	AuthExchange    string `json:"authExchange"`
	Environment     string `json:"environment"`
	Version         string `json:"version"`
	ResourceName    string `json:"resourceName"`
	UserSitesDomain string `json:"userSitesDomain"`
	LogResourceName string `json:"logResourceName"`
	SocialApiUri    string `json:"socialApiUri"`
	ApiUri          string `json:"apiUri"`
	MainUri         string `json:"mainUri"`
	SourceMapsUri   string `json:"sourceMapsUri"`
	Broker          struct {
		Uri string `json:"uri"`
	} `json:"broker"`
	AppsUri            string `json:"appsUri"`
	UploadsUri         string `json:"uploadsUri"`
	UploadsUriForGroup string `json:"uploadsUriForGroup"`
	FileFetchTimeout   int    `json:"fileFetchTimeout"`
	UserIdleMs         int    `json:"userIdleMs"`
	Embedly            struct {
		ApiKey string `json:"apiKey"`
	} `json:"embedly"`
	Github struct {
		ClientId string `json:"clientId"`
	} `json:"github"`
	Newkontrol struct {
		Url string `json:"url"`
	} `json:"newkontrol"`
	SessionCookie struct {
		MaxAge int  `json:"maxAge"`
		Secure bool `json:"secure"`
	} `json:"sessionCookie"`
	Stripe struct {
		Token string `json:"token"`
	} `json:"stripe"`
	ExternalProfiles struct {
		Google struct {
			Nicename string `json:"nicename"`
		} `json:"google"`
		Linkedin struct {
			Nicename string `json:"nicename"`
		} `json:"linkedin"`
		Twitter struct {
			Nicename string `json:"nicename"`
		} `json:"twitter"`
		Facebook struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"facebook"`
		Github struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"github"`
	} `json:"externalProfiles"`
	EntryPoint struct {
		Slug string `json:"slug"`
		Type string `json:"type"`
	} `json:"entryPoint"`
	Roles       []string      `json:"roles"`
	Permissions []interface{} `json:"permissions"`
	SiftScience string        `json:"siftScience"`
	Paypal      struct {
		FormUrl string `json:"formUrl"`
	} `json:"paypal"`
	Pubnub struct {
		SubscribeKey string `json:"subscribekey"`
		Enabled      bool   `json:"enabled"`
		SSL          bool   `json:"ssl"`
	} `json:"pubnub"`
	Collaboration struct {
		Timeout int `json:"timeout"`
	} `json:"collaboration"`
	PaymentBlockDuration float64 `json:"paymentBlockDuration"`
	DisabledFeatures     struct {
		Moderation bool `json:"moderation"`
		Teams      bool `json:"teams"`
		BotChannel bool `json:"botchannel"`
	} `json:"disabledFeatures"`
	ContentRotatorUrl string `json:"contentRotatorUrl"`
	Integration       struct {
		Url string `json:"url"`
	} `json:"integration"`
	WebhookMiddleware struct {
		Url string `json:"url"`
	} `json:"webhookMiddleware"`
	Google struct {
		ApiKey string `json:"apiKey"`
	} `json:"google"`
	Recaptcha struct {
		Key     string `json:"key"`
		Enabled bool   `json:"enabled"`
	} `json:"recaptcha"`
	Domains struct {
		Base string `json:"base"`
		Mail string `json:"mail"`
		Main string `json:"main"`
		Port string `json:"port"`
	} `json:"domains"`
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package engine

import (
	"context"
	"fmt"
	"testing"

	"vitess.io/vitess/go/test/utils"

	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/sqltypes"
)

func TestDistinct(t *testing.T) {
	type testCase struct {
		testName       string
		inputs         *sqltypes.Result
		expectedResult *sqltypes.Result
		expectedError  string
	}

	testCases := []*testCase{{
		testName:       "empty",
		inputs:         r("id1|col11|col12", "int64|varbinary|varbinary"),
		expectedResult: r("id1|col11|col12", "int64|varbinary|varbinary"),
	}, {
		testName:       "int64 numbers",
		inputs:         r("myid", "int64", "0", "1", "1", "null", "null"),
		expectedResult: r("myid", "int64", "0", "1", "null"),
	}, {
		testName:       "int64 numbers, two columns",
		inputs:         r("a|b", "int64|int64", "0|0", "1|1", "1|1", "null|null", "null|null", "1|2"),
		expectedResult: r("a|b", "int64|int64", "0|0", "1|1", "null|null", "1|2"),
	}, {
		testName:      "varchar columns",
		inputs:        r("myid", "varchar", "monkey", "horse"),
		expectedError: "types does not support hashcode yet",
	}}

	for _, tc := range testCases {
		t.Run(tc.testName+"-Execute", func(t *testing.T) {
			distinct := &Distinct{Source: &fakePrimitive{results: []*sqltypes.Result{tc.inputs}}}

			qr, err := distinct.Execute(&noopVCursor{ctx: context.Background()}, nil, true)
			if tc.expectedError == "" {
				require.NoError(t, err)
				got := fmt.Sprintf("%v", qr.Rows)
				expected := fmt.Sprintf("%v", tc.expectedResult.Rows)
				utils.MustMatch(t, expected, got, "result not what correct")
			} else {
				require.Error(t, err)
				require.Contains(t, err.Error(), tc.expectedError)
			}
		})
	}
}
<commit_msg>added test with hash collisions<commit_after>/*
Copyright 2020 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package engine

import (
	"context"
	"fmt"
	"testing"

	"vitess.io/vitess/go/test/utils"

	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/sqltypes"
)

func TestDistinct(t *testing.T) {
	type testCase struct {
		testName       string
		inputs         *sqltypes.Result
		expectedResult *sqltypes.Result
		expectedError  string
	}

	testCases := []*testCase{{
		testName:       "empty",
		inputs:         r("id1|col11|col12", "int64|varbinary|varbinary"),
		expectedResult: r("id1|col11|col12", "int64|varbinary|varbinary"),
	}, {
		testName:       "int64 numbers",
		inputs:         r("myid", "int64", "0", "1", "1", "null", "null"),
		expectedResult: r("myid", "int64", "0", "1", "null"),
	}, {
		testName:       "int64 numbers, two columns",
		inputs:         r("a|b", "int64|int64", "0|0", "1|1", "1|1", "null|null", "null|null", "1|2"),
		expectedResult: r("a|b", "int64|int64", "0|0", "1|1", "null|null", "1|2"),
	}, {
		testName:       "float64 columns designed to produce the same hashcode but not be equal",
		inputs:         r("a", "float64", "0.1", "0.2", "0.3", "0.4"),
		expectedResult: r("a", "float64", "0.1", "0.2", "0.3", "0.4"),
	}, {
		testName:      "varchar columns",
		inputs:        r("myid", "varchar", "monkey", "horse"),
		expectedError: "types does not support hashcode yet",
	}}

	for _, tc := range testCases {
		t.Run(tc.testName+"-Execute", func(t *testing.T) {
			distinct := &Distinct{Source: &fakePrimitive{results: []*sqltypes.Result{tc.inputs}}}

			qr, err := distinct.Execute(&noopVCursor{ctx: context.Background()}, nil, true)
			if tc.expectedError == "" {
				require.NoError(t, err)
				got := fmt.Sprintf("%v", qr.Rows)
				expected := fmt.Sprintf("%v", tc.expectedResult.Rows)
				utils.MustMatch(t, expected, got, "result not what correct")
			} else {
				require.Error(t, err)
				require.Contains(t, err.Error(), tc.expectedError)
			}
		})
	}
}
<|endoftext|>


<commit_before>package pstoremem

import (
	"sync"

	"github.com/libp2p/go-libp2p-core/peer"
	pstore "github.com/libp2p/go-libp2p-core/peerstore"
)

type memoryPeerMetadata struct {
	// store other data, like versions
	// ds ds.ThreadSafeDatastore
	ds       map[peer.ID]map[string]interface{}
	dslock   sync.RWMutex
	interned map[string]interface{}
}

var _ pstore.PeerMetadata = (*memoryPeerMetadata)(nil)

func NewPeerMetadata() *memoryPeerMetadata {
	return &memoryPeerMetadata{
		ds:       make(map[peer.ID]map[string]interface{}),
		interned: make(map[string]interface{}),
	}
}

func (ps *memoryPeerMetadata) Put(p peer.ID, key string, val interface{}) error {
	if err := p.Validate(); err != nil {
		return err
	}
	ps.dslock.Lock()
	defer ps.dslock.Unlock()
	if vals, ok := val.(string); ok && (key == "AgentVersion" || key == "ProtocolVersion") {
		if interned, ok := ps.interned[vals]; ok {
			val = interned
		} else {
			ps.interned[vals] = val
		}
	}
	m, ok := ps.ds[p]
	if !ok {
		m = make(map[string]interface{})
		ps.ds[p] = m
	}
	m[key] = val
	return nil
}

func (ps *memoryPeerMetadata) Get(p peer.ID, key string) (interface{}, error) {
	if err := p.Validate(); err != nil {
		return nil, err
	}
	ps.dslock.RLock()
	defer ps.dslock.RUnlock()
	m, ok := ps.ds[p]
	if !ok {
		return nil, pstore.ErrNotFound
	}
	val, ok := m[key]
	if !ok {
		return nil, pstore.ErrNotFound
	}
	return val, nil
}

func (ps *memoryPeerMetadata) RemovePeer(p peer.ID) {
	ps.dslock.Lock()
	delete(ps.ds, p)
	ps.dslock.Unlock()
}
<commit_msg>remove metadata interning<commit_after>package pstoremem

import (
	"sync"

	"github.com/libp2p/go-libp2p-core/peer"
	pstore "github.com/libp2p/go-libp2p-core/peerstore"
)

type memoryPeerMetadata struct {
	// store other data, like versions
	ds     map[peer.ID]map[string]interface{}
	dslock sync.RWMutex
}

var _ pstore.PeerMetadata = (*memoryPeerMetadata)(nil)

func NewPeerMetadata() *memoryPeerMetadata {
	return &memoryPeerMetadata{
		ds: make(map[peer.ID]map[string]interface{}),
	}
}

func (ps *memoryPeerMetadata) Put(p peer.ID, key string, val interface{}) error {
	if err := p.Validate(); err != nil {
		return err
	}
	ps.dslock.Lock()
	defer ps.dslock.Unlock()
	m, ok := ps.ds[p]
	if !ok {
		m = make(map[string]interface{})
		ps.ds[p] = m
	}
	m[key] = val
	return nil
}

func (ps *memoryPeerMetadata) Get(p peer.ID, key string) (interface{}, error) {
	if err := p.Validate(); err != nil {
		return nil, err
	}
	ps.dslock.RLock()
	defer ps.dslock.RUnlock()
	m, ok := ps.ds[p]
	if !ok {
		return nil, pstore.ErrNotFound
	}
	val, ok := m[key]
	if !ok {
		return nil, pstore.ErrNotFound
	}
	return val, nil
}

func (ps *memoryPeerMetadata) RemovePeer(p peer.ID) {
	ps.dslock.Lock()
	delete(ps.ds, p)
	ps.dslock.Unlock()
}
<|endoftext|>


<commit_before>package zkcli

import (
    "fmt"
    "strings"
    "time"

    "github.com/samuel/go-zookeeper/zk"
)

//回调：监视节点变化
type FuncWatchCallback func(path string, data []byte, err error)
//回调：监视子节点改变
type FuncWatchChildrenCallback func(path string, children []string, err error)

const (
    //zk session的超时时间，在此时间内，session可以自动保活
    //如果在此时间内，与特定server的连接断开，自动尝试重连其他服务器
    //重连成功之后，临时节点和watch依然有效
    DefaultConnectTimeout       = 5
)

type Conn struct {
    conn *zk.Conn
}
func (c *Conn) Conn() *zk.Conn {
    return c.conn
}
func (c *Conn) SetConn(conn *zk.Conn) {
    c.conn = conn
}

func (c *Conn) Close() {
    c.conn.Close()
}

func (c *Conn) Exists(path string) (bool, error) {
    exist, _, err := c.Conn().Exists(path)
    if err != nil {
        return false, err
    }
    return exist, nil
}

func (c *Conn) Get(path string) ([]byte, error) {
	data, _, err := c.Conn().Get(path);
	return data, err
}

func (c *Conn) GetChildren(path string) (map[string][]byte, error) {
	children := make(map[string][]byte)
	nodes, _, err := c.Conn().Children(path)
	if err != nil {
		return nil, err
	}
	for _, child := range nodes {
		data, err := c.Get(path + "/" + child)
		if err != nil {
			return nil, fmt.Errorf("child node %v get error %v", child, err)
		}
		children[child] = data
	}
	return children, nil
}

func (c *Conn) Write(path string, data []byte) error {
    exist, stat, err := c.conn.Exists(path)
    if err != nil {
        return err
    }
    if exist {
        _, err = c.conn.Set(path, data, stat.Version)
    } else {
        //不存在则创建
        // 永久节点
        flags := int32(0)
        _, err = doCreate(c.conn, path, data, flags)
    }
    return err
}

func (c *Conn) Watch(path string, cb FuncWatchCallback, stopCh chan struct{}) error {
    _, ch, err := getW(c.conn, path)
    if err != nil {
        return err
    }
    go func() {
        var data []byte
        for {
            select {
            case <-stopCh:
                return
            case ev := <-ch:
                if ev.Err != nil {
                    //错误回调
                    cb(path, nil, ev.Err)
                    return
                }
                if ev.Path != path {
                    cb(path, nil, fmt.Errorf("mismatched path %v %v", ev.Path, path))
                    return
                }
            }
            // 获取变化后的节点数据
            // 并更新watcher（zookeeper的watcher是一次性的）
            data, ch, err = getW(c.conn, path)
            if err != nil {
                //错误回调
                cb(path, nil, err)
                return
            }
            //数据回调
            cb(path, data, nil)
        }
    }()

    return nil
}

func (c *Conn) Create(path string, data []byte) error {
    exist, _, err := c.Conn().Exists(path)
    if err != nil {
        return err
    }
    if exist {
        return nil
    }
    _, err = doCreate(c.Conn(), path, data, 0)
    return err
}

func (c *Conn) CreateEphemeral(path string, data []byte) error {
    exist, _, err := c.Conn().Exists(path)
    if err != nil {
        return err
    }
    if exist {
        return nil
    }
    // 临时节点
    _, err = doCreate(c.Conn(), path, data, int32(zk.FlagEphemeral))
    return err
}

func (c *Conn) CreateSequence(path string, data []byte) (string, error) {
    flags := int32(zk.FlagSequence | zk.FlagEphemeral)
    return doCreate(c.Conn(), path, data, flags)
}

func (c *Conn) WatchChildren(path string, cb FuncWatchChildrenCallback, stopCh chan struct{}) error {
    _, ch, err := childrenW(c.Conn(), path)
    if err != nil {
        return err
    }
    go func() {
        for {
            select {
            case <-stopCh:
                return
            case ev := <-ch:
                if ev.Err != nil {
                    //错误回调
                    cb(path, nil, ev.Err)
                    return
                }
                if ev.Path != path {
                    cb(path, nil, fmt.Errorf("mismatched path %v %v", ev.Path, path))
                    return
                }
            }
            // 获取变化后的节点数据
            // 并更新watcher（zookeeper的watcher是一次性的）
            var children []string
            children, ch, err = childrenW(c.Conn(), path)
            if err != nil {
                //错误回调
                cb(path, nil, err)
                return
            }
            //数据回调
            cb(path, children, nil)
        }
    }()

    return nil
}

func Connect(zkAddr string) (*Conn, error) {
    if len(zkAddr) == 0 {
        return nil, fmt.Errorf("empty zkAddr")
    }
    zks := strings.Split(zkAddr, ",")
    conn, _, err := zk.Connect(zks, time.Second * DefaultConnectTimeout)
    if err != nil {
        return nil, fmt.Errorf("err connect to zk<%v>: %v", zkAddr, err)
    }

    c := new(Conn)
    c.SetConn(conn)
    return c, nil
}

////////////////////////////////////////////////////////////////////////////

func getW(zkConn *zk.Conn, path string) ([]byte, <-chan zk.Event, error) {
    data, _, ch, err := zkConn.GetW(path)
    return data, ch, err
}
func childrenW(zkConn *zk.Conn, path string) ([]string, <-chan zk.Event, error) {
    children, _, ch, err := zkConn.ChildrenW(path)
    return children, ch, err
}

func doCreate(zkConn *zk.Conn, path string, data []byte, flags int32) (string, error) {
    //TODO 权限控制
    acl := zk.WorldACL(zk.PermAll)
    path, err := zkConn.Create(path, data, flags, acl)
    if err != nil {
        return "", err
    }

    return path, nil
}

<commit_msg>enhance zkcli functions<commit_after>package zkcli

import (
    "fmt"
    "strings"
    "time"

    "github.com/samuel/go-zookeeper/zk"
)

//回调：监视节点变化
type FuncWatchCallback func(path string, data []byte, err error)
//回调：监视子节点改变
type FuncWatchChildrenCallback func(path string, children []string, err error)

const (
    //zk session的超时时间，在此时间内，session可以自动保活
    //如果在此时间内，与特定server的连接断开，自动尝试重连其他服务器
    //重连成功之后，临时节点和watch依然有效
    DefaultConnectTimeout       = 5
)

type Conn struct {
    conn *zk.Conn
}
func (c *Conn) Conn() *zk.Conn {
    return c.conn
}
func (c *Conn) SetConn(conn *zk.Conn) {
    c.conn = conn
}

func (c *Conn) Close() {
    c.conn.Close()
}

func (c *Conn) Exists(path string) (bool, error) {
    exist, _, err := c.Conn().Exists(path)
    if err != nil {
        return false, err
    }
    return exist, nil
}

func (c *Conn) Get(path string) ([]byte, error) {
	data, _, err := c.Conn().Get(path);
	return data, err
}

func (c *Conn) GetChildren(path string) (map[string][]byte, error) {
	children := make(map[string][]byte)
	nodes, _, err := c.Conn().Children(path)
	if err != nil {
		return nil, err
	}
	for _, child := range nodes {
		data, err := c.Get(path + "/" + child)
		if err != nil {
			return nil, fmt.Errorf("child node %v get error %v", child, err)
		}
		children[child] = data
	}
	return children, nil
}

func (c *Conn) Write(path string, data []byte) error {
	//永久节点
	return c.write(path, data, int32(0))
}

func (c *Conn) WriteEphemeral(path string, data []byte) error {
	//临时节点
	return c.write(path, data, int32(zk.FlagEphemeral))
}

func (c *Conn) write(path string, data []byte, flags int32) error {
    exist, stat, err := c.conn.Exists(path)
    if err != nil {
        return err
    }
    if exist {
        _, err = c.conn.Set(path, data, stat.Version)
    } else {
        //不存在则创建
        _, err = doCreate(c.conn, path, data, flags)
    }
    return err
}

func (c *Conn) Watch(path string, cb FuncWatchCallback, stopCh chan struct{}) error {
    _, ch, err := getW(c.conn, path)
    if err != nil {
        return err
    }
    go func() {
        var data []byte
        for {
            select {
            case <-stopCh:
                return
            case ev := <-ch:
                if ev.Err != nil {
                    //错误回调
                    cb(path, nil, ev.Err)
                    return
                }
                if ev.Path != path {
                    cb(path, nil, fmt.Errorf("mismatched path %v %v", ev.Path, path))
                    return
                }
            }
            // 获取变化后的节点数据
            // 并更新watcher（zookeeper的watcher是一次性的）
            data, ch, err = getW(c.conn, path)
            if err != nil {
                //错误回调
                cb(path, nil, err)
                return
            }
            //数据回调
            cb(path, data, nil)
        }
    }()

    return nil
}

func (c *Conn) Create(path string, data []byte) error {
    exist, _, err := c.Conn().Exists(path)
    if err != nil {
        return err
    }
    if exist {
        return nil
    }
    _, err = doCreate(c.Conn(), path, data, 0)
    return err
}

func (c *Conn) CreateEphemeral(path string, data []byte) error {
    exist, _, err := c.Conn().Exists(path)
    if err != nil {
        return err
    }
    if exist {
        return nil
    }
    // 临时节点
    _, err = doCreate(c.Conn(), path, data, int32(zk.FlagEphemeral))
    return err
}

func (c *Conn) CreateSequence(path string, data []byte) (string, error) {
    flags := int32(zk.FlagSequence | zk.FlagEphemeral)
    return doCreate(c.Conn(), path, data, flags)
}

func (c *Conn) WatchChildren(path string, cb FuncWatchChildrenCallback, stopCh chan struct{}) error {
    _, ch, err := childrenW(c.Conn(), path)
    if err != nil {
        return err
    }
    go func() {
        for {
            select {
            case <-stopCh:
                return
            case ev := <-ch:
                if ev.Err != nil {
                    //错误回调
                    cb(path, nil, ev.Err)
                    return
                }
                if ev.Path != path {
                    cb(path, nil, fmt.Errorf("mismatched path %v %v", ev.Path, path))
                    return
                }
            }
            // 获取变化后的节点数据
            // 并更新watcher（zookeeper的watcher是一次性的）
            var children []string
            children, ch, err = childrenW(c.Conn(), path)
            if err != nil {
                //错误回调
                cb(path, nil, err)
                return
            }
            //数据回调
            cb(path, children, nil)
        }
    }()

    return nil
}

func Connect(zkAddr string) (*Conn, error) {
    if len(zkAddr) == 0 {
        return nil, fmt.Errorf("empty zkAddr")
    }
    zks := strings.Split(zkAddr, ",")
    conn, _, err := zk.Connect(zks, time.Second * DefaultConnectTimeout)
    if err != nil {
        return nil, fmt.Errorf("err connect to zk<%v>: %v", zkAddr, err)
    }

    c := new(Conn)
    c.SetConn(conn)
    return c, nil
}

////////////////////////////////////////////////////////////////////////////

func getW(zkConn *zk.Conn, path string) ([]byte, <-chan zk.Event, error) {
    data, _, ch, err := zkConn.GetW(path)
    return data, ch, err
}
func childrenW(zkConn *zk.Conn, path string) ([]string, <-chan zk.Event, error) {
    children, _, ch, err := zkConn.ChildrenW(path)
    return children, ch, err
}

func doCreate(zkConn *zk.Conn, path string, data []byte, flags int32) (string, error) {
    //TODO 权限控制
    acl := zk.WorldACL(zk.PermAll)
    path, err := zkConn.Create(path, data, flags, acl)
    if err != nil {
        return "", err
    }

    return path, nil
}

<|endoftext|>