<commit_before>// Copyright 2018 Google LLC
//
// 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
//
//     https://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 gengapic

import (
	"fmt"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/ptypes/duration"
	conf "github.com/googleapis/gapic-generator-go/internal/grpc_service_config"
	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
	"google.golang.org/genproto/googleapis/api/annotations"
	"google.golang.org/genproto/googleapis/rpc/code"
)

func (g *generator) clientOptions(serv *descriptor.ServiceDescriptorProto, servName string) error {
	p := g.printf

	// CallOptions struct
	{
		p("// %[1]sCallOptions contains the retry settings for each method of %[1]sClient.", servName)
		p("type %sCallOptions struct {", servName)
		for _, m := range serv.Method {
			p("%s []gax.CallOption", *m.Name)
		}
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{"gax", "github.com/googleapis/gax-go/v2"}] = true
	}

	// defaultClientOptions
	{
		var host string
		if eHost, err := proto.GetExtension(serv.Options, annotations.E_DefaultHost); err == nil {
			host = *eHost.(*string)
		} else {
			fqn := g.descInfo.ParentFile[serv].GetPackage() + "." + serv.GetName()
			return fmt.Errorf("service %q is missing option google.api.default_host", fqn)
		}

		if !strings.Contains(host, ":") {
			host += ":443"
		}

		p("func default%sClientOptions() []option.ClientOption {", servName)
		p("  return []option.ClientOption{")
		p("    option.WithEndpoint(%q),", host)
		p("    option.WithGRPCDialOption(grpc.WithDisableServiceConfig()),")
		p("    option.WithScopes(DefaultAuthScopes()...),")
		p("    option.WithGRPCDialOption(grpc.WithDefaultCallOptions(")
		p("      grpc.MaxCallRecvMsgSize(math.MaxInt32))),")
		p("  }")
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{Path: "math"}] = true
		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/api/option"}] = true
	}

	// defaultCallOptions
	{
		sFQN := fmt.Sprintf("%s.%s", g.descInfo.ParentFile[serv].GetPackage(), serv.GetName())
		policies := map[string]*conf.MethodConfig_RetryPolicy{}
		reqLimits := map[string]int{}
		resLimits := map[string]int{}

		var methCfgs []*conf.MethodConfig
		if g.grpcConf != nil {
			methCfgs = g.grpcConf.GetMethodConfig()
		}

		// gather retry policies from MethodConfigs
		for _, mc := range methCfgs {
			for _, name := range mc.GetName() {
				base := name.GetService()

				// skip the Name entry if it's not the current service
				if base != sFQN {
					continue
				}

				// individual method config, overwrites service-level config
				if name.GetMethod() != "" {
					base = base + "." + name.GetMethod()
					policies[base] = mc.GetRetryPolicy()

					if maxReq := mc.GetMaxRequestMessageBytes(); maxReq != nil {
						reqLimits[base] = int(maxReq.GetValue())
					}

					if maxRes := mc.GetMaxResponseMessageBytes(); maxRes != nil {
						resLimits[base] = int(maxRes.GetValue())
					}

					continue
				}

				// service-level config, apply to all *unset* methods
				for _, m := range serv.GetMethod() {
					// build fully-qualified name
					fqn := base + "." + m.GetName()

					// set retry config
					if _, ok := policies[fqn]; !ok {
						policies[fqn] = mc.GetRetryPolicy()
					}

					// set max request size limit
					if maxReq := mc.GetMaxRequestMessageBytes(); maxReq != nil {
						if _, ok := reqLimits[fqn]; !ok {
							reqLimits[fqn] = int(maxReq.GetValue())
						}
					}

					// set max response size limit
					if maxRes := mc.GetMaxResponseMessageBytes(); maxRes != nil {
						if _, ok := resLimits[fqn]; !ok {
							resLimits[fqn] = int(maxRes.GetValue())
						}
					}
				}
			}
		}

		if len(policies) > 0 {
			g.imports[pbinfo.ImportSpec{Path: "time"}] = true
			g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc/codes"}] = true
		}

		// read retry params from gRPC ServiceConfig
		p("func default%[1]sCallOptions() *%[1]sCallOptions {", servName)
		p("  return &%sCallOptions{", servName)
		for _, m := range serv.GetMethod() {
			mFQN := sFQN + "." + m.GetName()
			p("%s: []gax.CallOption{", m.GetName())

			if maxReq, ok := reqLimits[mFQN]; ok {
				p("gax.WithGRPCOptions(grpc.MaxCallSendMsgSize(%d)),", maxReq)
			}

			if maxRes, ok := resLimits[mFQN]; ok {
				p("gax.WithGRPCOptions(grpc.MaxCallRecvMsgSize(%d)),", maxRes)
			}

			if rp, ok := policies[mFQN]; ok && rp != nil {
				p("gax.WithRetry(func() gax.Retryer {")
				p("  return gax.OnCodes([]codes.Code{")
				for _, c := range rp.GetRetryableStatusCodes() {
					cstr := c.String()

					// Go uses the American-English spelling with a single "L"
					if c == code.Code_CANCELLED {
						cstr = "Canceled"
					}

					p("    codes.%s,", snakeToCamel(cstr))
				}
				p("	 }, gax.Backoff{")
				// this ignores max_attempts
				p("		Initial:    %d * time.Millisecond,", durationToMillis(rp.GetInitialBackoff()))
				p("		Max:        %d * time.Millisecond,", durationToMillis(rp.GetMaxBackoff()))
				p("		Multiplier: %.2f,", rp.GetBackoffMultiplier())
				p("	 })")
				p("}),")
			}
			p("},")
		}
		p("  }")
		p("}")
		p("")
	}

	return nil
}

func durationToMillis(d *duration.Duration) int64 {
	return d.GetSeconds()*1000 + int64(d.GetNanos()/1000000)
}

func (g *generator) clientInit(serv *descriptor.ServiceDescriptorProto, servName string) error {
	p := g.printf

	var hasLRO bool
	for _, m := range serv.Method {
		if g.isLRO(m) {
			hasLRO = true
			break
		}
	}

	imp, err := g.descInfo.ImportSpec(serv)
	if err != nil {
		return err
	}

	// client struct
	{
		p("// %sClient is a client for interacting with %s.", servName, g.apiName)
		p("//")
		p("// Methods, except Close, may be called concurrently. However, fields must not be modified concurrently with method calls.")
		p("type %sClient struct {", servName)

		p("// Connection pool of gRPC connections to the service.")
		p("connPool gtransport.ConnPool")
		p("")

		p("// The gRPC API client.")
		p("%s %s.%sClient", grpcClientField(servName), imp.Name, serv.GetName())
		p("")

		if hasLRO {
			p("// LROClient is used internally to handle longrunning operations.")
			p("// It is exposed so that its CallOptions can be modified if required.")
			p("// Users should not Close this client.")
			p("LROClient *lroauto.OperationsClient")
			p("")

			g.imports[pbinfo.ImportSpec{Name: "lroauto", Path: "cloud.google.com/go/longrunning/autogen"}] = true
		}

		p("// The call options for this service.")
		p("CallOptions *%sCallOptions", servName)
		p("")

		p("// The x-goog-* metadata to be sent with each request.")
		p("xGoogMetadata metadata.MD")
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc"}] = true
		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc/metadata"}] = true
	}

	// Client constructor
	{
		clientName := camelToSnake(serv.GetName())
		clientName = strings.Replace(clientName, "_", " ", -1)

		p("// New%sClient creates a new %s client.", servName, clientName)
		p("//")
		g.comment(g.comments[serv])
		p("func New%[1]sClient(ctx context.Context, opts ...option.ClientOption) (*%[1]sClient, error) {", servName)
		p("  connPool, err := gtransport.DialPool(ctx, append(default%sClientOptions(), opts...)...)", servName)
		p("  if err != nil {")
		p("    return nil, err")
		p("  }")
		p("  c := &%sClient{", servName)
		p("    connPool:    connPool,")
		p("    CallOptions: default%sCallOptions(),", servName)
		p("")
		p("    %s: %s.New%sClient(connPool),", grpcClientField(servName), imp.Name, serv.GetName())
		p("  }")
		p("  c.setGoogleClientInfo()")
		p("")

		if hasLRO {
			p("  c.LROClient, err = lroauto.NewOperationsClient(ctx, gtransport.WithConnPool(connPool))")
			p("  if err != nil {")
			p("    // This error \"should not happen\", since we are just reusing old connection pool")
			p("    // and never actually need to dial.")
			p("    // If this does happen, we could leak connp. However, we cannot close conn:")
			p("    // If the user invoked the constructor with option.WithGRPCConn,")
			p("    // we would close a connection that's still in use.")
			p("    // TODO: investigate error conditions.")
			p("    return nil, err")
			p("  }")
		}

		p("  return c, nil")
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{Name: "gtransport", Path: "google.golang.org/api/transport/grpc"}] = true
		g.imports[pbinfo.ImportSpec{Path: "context"}] = true
	}

	// Connection()
	{
		p("// Connection returns a connection to the API service.")
		p("//")
		p("// Deprecated.")
		p("func (c *%sClient) Connection() *grpc.ClientConn {", servName)
		p("  return c.connPool.Conn()")
		p("}")
		p("")
	}

	// Close()
	{
		p("// Close closes the connection to the API service. The user should invoke this when")
		p("// the client is no longer required.")
		p("func (c *%sClient) Close() error {", servName)
		p("  return c.connPool.Close()")
		p("}")
		p("")
	}

	// setGoogleClientInfo
	{
		p("// setGoogleClientInfo sets the name and version of the application in")
		p("// the `x-goog-api-client` header passed on each request. Intended for")
		p("// use by Google-written clients.")
		p("func (c *%sClient) setGoogleClientInfo(keyval ...string) {", servName)
		p(`  kv := append([]string{"gl-go", versionGo()}, keyval...)`)
		p(`  kv = append(kv, "gapic", versionClient, "gax", gax.Version, "grpc", grpc.Version)`)
		p(`  c.xGoogMetadata = metadata.Pairs("x-goog-api-client", gax.XGoogHeader(kv...))`)
		p("}")
		p("")
	}

	return nil
}
<commit_msg>gapic: fix naive retry related imports (#317)<commit_after>// Copyright 2018 Google LLC
//
// 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
//
//     https://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 gengapic

import (
	"fmt"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/ptypes/duration"
	conf "github.com/googleapis/gapic-generator-go/internal/grpc_service_config"
	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
	"google.golang.org/genproto/googleapis/api/annotations"
	"google.golang.org/genproto/googleapis/rpc/code"
)

func (g *generator) clientOptions(serv *descriptor.ServiceDescriptorProto, servName string) error {
	p := g.printf

	// CallOptions struct
	{
		p("// %[1]sCallOptions contains the retry settings for each method of %[1]sClient.", servName)
		p("type %sCallOptions struct {", servName)
		for _, m := range serv.Method {
			p("%s []gax.CallOption", *m.Name)
		}
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{"gax", "github.com/googleapis/gax-go/v2"}] = true
	}

	// defaultClientOptions
	{
		var host string
		if eHost, err := proto.GetExtension(serv.Options, annotations.E_DefaultHost); err == nil {
			host = *eHost.(*string)
		} else {
			fqn := g.descInfo.ParentFile[serv].GetPackage() + "." + serv.GetName()
			return fmt.Errorf("service %q is missing option google.api.default_host", fqn)
		}

		if !strings.Contains(host, ":") {
			host += ":443"
		}

		p("func default%sClientOptions() []option.ClientOption {", servName)
		p("  return []option.ClientOption{")
		p("    option.WithEndpoint(%q),", host)
		p("    option.WithGRPCDialOption(grpc.WithDisableServiceConfig()),")
		p("    option.WithScopes(DefaultAuthScopes()...),")
		p("    option.WithGRPCDialOption(grpc.WithDefaultCallOptions(")
		p("      grpc.MaxCallRecvMsgSize(math.MaxInt32))),")
		p("  }")
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{Path: "math"}] = true
		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/api/option"}] = true
	}

	// defaultCallOptions
	{
		sFQN := fmt.Sprintf("%s.%s", g.descInfo.ParentFile[serv].GetPackage(), serv.GetName())
		policies := map[string]*conf.MethodConfig_RetryPolicy{}
		reqLimits := map[string]int{}
		resLimits := map[string]int{}

		var methCfgs []*conf.MethodConfig
		if g.grpcConf != nil {
			methCfgs = g.grpcConf.GetMethodConfig()
		}

		// gather retry policies from MethodConfigs
		for _, mc := range methCfgs {
			for _, name := range mc.GetName() {
				base := name.GetService()

				// skip the Name entry if it's not the current service
				if base != sFQN {
					continue
				}

				// individual method config, overwrites service-level config
				if name.GetMethod() != "" {
					base = base + "." + name.GetMethod()
					policies[base] = mc.GetRetryPolicy()

					if maxReq := mc.GetMaxRequestMessageBytes(); maxReq != nil {
						reqLimits[base] = int(maxReq.GetValue())
					}

					if maxRes := mc.GetMaxResponseMessageBytes(); maxRes != nil {
						resLimits[base] = int(maxRes.GetValue())
					}

					continue
				}

				// service-level config, apply to all *unset* methods
				for _, m := range serv.GetMethod() {
					// build fully-qualified name
					fqn := base + "." + m.GetName()

					// set retry config
					if _, ok := policies[fqn]; !ok {
						policies[fqn] = mc.GetRetryPolicy()
					}

					// set max request size limit
					if maxReq := mc.GetMaxRequestMessageBytes(); maxReq != nil {
						if _, ok := reqLimits[fqn]; !ok {
							reqLimits[fqn] = int(maxReq.GetValue())
						}
					}

					// set max response size limit
					if maxRes := mc.GetMaxResponseMessageBytes(); maxRes != nil {
						if _, ok := resLimits[fqn]; !ok {
							resLimits[fqn] = int(maxRes.GetValue())
						}
					}
				}
			}
		}

		// read retry params from gRPC ServiceConfig
		p("func default%[1]sCallOptions() *%[1]sCallOptions {", servName)
		p("  return &%sCallOptions{", servName)
		for _, m := range serv.GetMethod() {
			mFQN := sFQN + "." + m.GetName()
			p("%s: []gax.CallOption{", m.GetName())

			if maxReq, ok := reqLimits[mFQN]; ok {
				p("gax.WithGRPCOptions(grpc.MaxCallSendMsgSize(%d)),", maxReq)
			}

			if maxRes, ok := resLimits[mFQN]; ok {
				p("gax.WithGRPCOptions(grpc.MaxCallRecvMsgSize(%d)),", maxRes)
			}

			if rp, ok := policies[mFQN]; ok && rp != nil {
				p("gax.WithRetry(func() gax.Retryer {")
				p("  return gax.OnCodes([]codes.Code{")
				for _, c := range rp.GetRetryableStatusCodes() {
					cstr := c.String()

					// Go uses the American-English spelling with a single "L"
					if c == code.Code_CANCELLED {
						cstr = "Canceled"
					}

					p("    codes.%s,", snakeToCamel(cstr))
				}
				p("	 }, gax.Backoff{")
				// this ignores max_attempts
				p("		Initial:    %d * time.Millisecond,", durationToMillis(rp.GetInitialBackoff()))
				p("		Max:        %d * time.Millisecond,", durationToMillis(rp.GetMaxBackoff()))
				p("		Multiplier: %.2f,", rp.GetBackoffMultiplier())
				p("	 })")
				p("}),")

				// include imports necessary for retry configuration
				g.imports[pbinfo.ImportSpec{Path: "time"}] = true
				g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc/codes"}] = true
			}
			p("},")
		}
		p("  }")
		p("}")
		p("")
	}

	return nil
}

func durationToMillis(d *duration.Duration) int64 {
	return d.GetSeconds()*1000 + int64(d.GetNanos()/1000000)
}

func (g *generator) clientInit(serv *descriptor.ServiceDescriptorProto, servName string) error {
	p := g.printf

	var hasLRO bool
	for _, m := range serv.Method {
		if g.isLRO(m) {
			hasLRO = true
			break
		}
	}

	imp, err := g.descInfo.ImportSpec(serv)
	if err != nil {
		return err
	}

	// client struct
	{
		p("// %sClient is a client for interacting with %s.", servName, g.apiName)
		p("//")
		p("// Methods, except Close, may be called concurrently. However, fields must not be modified concurrently with method calls.")
		p("type %sClient struct {", servName)

		p("// Connection pool of gRPC connections to the service.")
		p("connPool gtransport.ConnPool")
		p("")

		p("// The gRPC API client.")
		p("%s %s.%sClient", grpcClientField(servName), imp.Name, serv.GetName())
		p("")

		if hasLRO {
			p("// LROClient is used internally to handle longrunning operations.")
			p("// It is exposed so that its CallOptions can be modified if required.")
			p("// Users should not Close this client.")
			p("LROClient *lroauto.OperationsClient")
			p("")

			g.imports[pbinfo.ImportSpec{Name: "lroauto", Path: "cloud.google.com/go/longrunning/autogen"}] = true
		}

		p("// The call options for this service.")
		p("CallOptions *%sCallOptions", servName)
		p("")

		p("// The x-goog-* metadata to be sent with each request.")
		p("xGoogMetadata metadata.MD")
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc"}] = true
		g.imports[pbinfo.ImportSpec{Path: "google.golang.org/grpc/metadata"}] = true
	}

	// Client constructor
	{
		clientName := camelToSnake(serv.GetName())
		clientName = strings.Replace(clientName, "_", " ", -1)

		p("// New%sClient creates a new %s client.", servName, clientName)
		p("//")
		g.comment(g.comments[serv])
		p("func New%[1]sClient(ctx context.Context, opts ...option.ClientOption) (*%[1]sClient, error) {", servName)
		p("  connPool, err := gtransport.DialPool(ctx, append(default%sClientOptions(), opts...)...)", servName)
		p("  if err != nil {")
		p("    return nil, err")
		p("  }")
		p("  c := &%sClient{", servName)
		p("    connPool:    connPool,")
		p("    CallOptions: default%sCallOptions(),", servName)
		p("")
		p("    %s: %s.New%sClient(connPool),", grpcClientField(servName), imp.Name, serv.GetName())
		p("  }")
		p("  c.setGoogleClientInfo()")
		p("")

		if hasLRO {
			p("  c.LROClient, err = lroauto.NewOperationsClient(ctx, gtransport.WithConnPool(connPool))")
			p("  if err != nil {")
			p("    // This error \"should not happen\", since we are just reusing old connection pool")
			p("    // and never actually need to dial.")
			p("    // If this does happen, we could leak connp. However, we cannot close conn:")
			p("    // If the user invoked the constructor with option.WithGRPCConn,")
			p("    // we would close a connection that's still in use.")
			p("    // TODO: investigate error conditions.")
			p("    return nil, err")
			p("  }")
		}

		p("  return c, nil")
		p("}")
		p("")

		g.imports[pbinfo.ImportSpec{Name: "gtransport", Path: "google.golang.org/api/transport/grpc"}] = true
		g.imports[pbinfo.ImportSpec{Path: "context"}] = true
	}

	// Connection()
	{
		p("// Connection returns a connection to the API service.")
		p("//")
		p("// Deprecated.")
		p("func (c *%sClient) Connection() *grpc.ClientConn {", servName)
		p("  return c.connPool.Conn()")
		p("}")
		p("")
	}

	// Close()
	{
		p("// Close closes the connection to the API service. The user should invoke this when")
		p("// the client is no longer required.")
		p("func (c *%sClient) Close() error {", servName)
		p("  return c.connPool.Close()")
		p("}")
		p("")
	}

	// setGoogleClientInfo
	{
		p("// setGoogleClientInfo sets the name and version of the application in")
		p("// the `x-goog-api-client` header passed on each request. Intended for")
		p("// use by Google-written clients.")
		p("func (c *%sClient) setGoogleClientInfo(keyval ...string) {", servName)
		p(`  kv := append([]string{"gl-go", versionGo()}, keyval...)`)
		p(`  kv = append(kv, "gapic", versionClient, "gax", gax.Version, "grpc", grpc.Version)`)
		p(`  c.xGoogMetadata = metadata.Pairs("x-goog-api-client", gax.XGoogHeader(kv...))`)
		p("}")
		p("")
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"testing"

	"golang.org/x/net/context"

	. "github.com/jacobsa/ogletest"
)

func TestBlobStore(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type VisitorTest struct {
	ctx context.Context
}

func init() { RegisterTestSuite(&VisitorTest{}) }

func (t *VisitorTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *VisitorTest) ScoresAlreadyPresent_Empty() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) ScoresAlreadyPresent_NonEmpty() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) Symlink() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) Directory() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) File_Empty() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) File_LastChunkIsFull() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) File_LastChunkIsPartial() {
	AssertTrue(false, "TODO")
}
<commit_msg>VisitorTest.ScoresAlreadyPresent_NonEmpty<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/blob"
	"github.com/jacobsa/comeback/internal/blob/mock"
	. "github.com/jacobsa/ogletest"
)

func TestBlobStore(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type VisitorTest struct {
	ctx       context.Context
	blobStore blob.Store

	node fsNode
}

func init() { RegisterTestSuite(&VisitorTest{}) }

func (t *VisitorTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.blobStore = mock_blob.NewMockStore(ti.MockController, "blobStore")
}

func (t *VisitorTest) call() (err error) {
	visitor := newVisitor(t.blobStore, make(chan *fsNode, 1))
	err = visitor.Visit(t.ctx, &t.node)
	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *VisitorTest) ScoresAlreadyPresent_Empty() {
	scores := []blob.Score{}
	t.node.Scores = scores

	err := t.call()
	AssertEq(nil, err)
	ExpectEq(scores, t.node.Scores)
}

func (t *VisitorTest) ScoresAlreadyPresent_NonEmpty() {
	scores := []blob.Score{blob.ComputeScore([]byte("taco"))}
	t.node.Scores = scores

	err := t.call()
	AssertEq(nil, err)
	ExpectEq(scores, t.node.Scores)
}

func (t *VisitorTest) Symlink() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) Directory() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) File_Empty() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) File_LastChunkIsFull() {
	AssertTrue(false, "TODO")
}

func (t *VisitorTest) File_LastChunkIsPartial() {
	AssertTrue(false, "TODO")
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The ContainerOps Authors All rights reserved.

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 models

import (
	"encoding/json"
	"time"
)

//DockerV1 is Docker Repository V1 repository.
type DockerV1 struct {
	ID          int64      `json:"id" gorm:"primary_key"`
	Namespace   string     `json:"namespace" sql:"not null;type:varchar(255)" gorm:"unique_index:dockerv1_repository"`
	Repository  string     `json:"repository" sql:"not null;type:varchar(255)" gorm:"unique_index:dockerv1_repository"`
	JSON        string     `json:"json" sql:"null;type:text"`
	Manifests   string     `json:"manifests" sql:"null;type:text"`
	Agent       string     `json:"agent" sql:"null;type:text"`
	Description string     `json:"description" sql:"null;type:text"`
	Size        int64      `json:"size" sql:"default:0"`
	Locked      bool       `json:"locked" sql:"default:false"` //When create/update the repository, the locked will be true.
	CreatedAt   time.Time  `json:"create_at" sql:""`
	UpdatedAt   time.Time  `json:"update_at" sql:""`
	DeletedAt   *time.Time `json:"delete_at" sql:"index"`
}

//TableName in mysql is "docker_v1".
func (r *DockerV1) TableName() string {
	return "docker_v1"
}

//DockerImageV1 is
type DockerImageV1 struct {
	ID         int64      `json:"id" gorm:"primary_key"`
	ImageID    string     `json:"image_id" sql:"not null;unique;varchar(255)"`
	JSON       string     `json:"json" sql:"null;type:text"`
	Ancestry   string     `json:"ancestry" sql:"null;type:text"`
	Checksum   string     `json:"checksum" sql:"null;type:varchar(255)"`
	Payload    string     `json:"payload" sql:"null;type:varchar(255)"`
	Path       string     `json:"path" sql:"null;type:text"`
	OSS        string     `json:"oss" sql:"null;type:text"`
	Size       int64      `json:"size" sql:"default:0"`
	Uploaded   bool       `json:"uploaded" sql:"default:false"`
	Checksumed bool       `json:"checksumed" sql:"default:false"`
	Locked     bool       `json:"locked" sql:"default:false"`
	CreatedAt  time.Time  `json:"create_at" sql:""`
	UpdatedAt  time.Time  `json:"update_at" sql:""`
	DeletedAt  *time.Time `json:"delete_at" sql:"index"`
}

//TableName in mysql is "docker_image_v1".
func (i *DockerImageV1) TableName() string {
	return "docker_image_v1"
}

//DockerTagV1 is
type DockerTagV1 struct {
	ID        int64      `json:"id" gorm:"primary_key"`
	DockerV1  int64      `json:"docker_v1" sql:"not null;default:0"`
	Tag       string     `json:"tag" sql:"not null;varchar(255)"`
	ImageID   string     `json:"image_id" sql:"not null;varchar(255)"`
	CreatedAt time.Time  `json:"create_at" sql:""`
	UpdatedAt time.Time  `json:"update_at" sql:""`
	DeletedAt *time.Time `json:"delete_at" sql:"index"`
}

//TableName in mysql is "docker_tag_v1".
func (t *DockerTagV1) TableName() string {
	return "docker_tag_v1"
}

//Put function will create or update repository.
func (r *DockerV1) Put(namespace, repository, json, agent string) error {
	r.Namespace, r.Repository, r.JSON, r.Agent, r.Locked = namespace, repository, json, agent, true

	tx := db.Begin()

	if err := tx.Debug().Where("namespace = ? AND repository = ? ", namespace, repository).FirstOrCreate(&r).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&r).Updates(map[string]interface{}{"json": json, "agent": agent, "locked": true}).Error; err != nil {
		tx.Rollback()
		return err
	} else if err == nil {
		tx.Commit()
		return nil
	}

	tx.Commit()
	return nil
}

//Unlocked is Unlocked repository data so could pull.
func (r *DockerV1) Unlocked(namespace, repository string) error {
	tx := db.Begin()

	if err := tx.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&r).Updates(map[string]interface{}{"locked": false}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}

//Get return Docker V1 repository data.
func (r *DockerV1) Get(namespace, repository string) (DockerV1, error) {
	if err := db.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		return *new(DockerV1), err
	} else {
		return *r, nil
	}
}

//GetTags return tas data of repository.
func (r *DockerV1) GetTags(namespace, repository string) (map[string]string, error) {
	if err := db.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		return map[string]string{}, err
	} else {
		var tags []DockerTagV1
		result := map[string]string{}

		if err := db.Debug().Where("docker_v1 = ?", r.ID).Find(&tags).Error; err != nil {
			return map[string]string{}, err
		}

		for _, tag := range tags {
			result[tag.Tag] = tag.ImageID
		}

		return result, nil
	}
}

//Get is search image by ImageID.
func (i *DockerImageV1) Get(imageID string) (DockerImageV1, error) {
	if err := db.Debug().Where("image_id = ?", imageID).First(&i).Error; err != nil {
		return *i, err
	} else {
		return *i, nil
	}
}

//PutJSON is put image json by ImageID.
func (i *DockerImageV1) PutJSON(imageID, json string) error {
	i.ImageID = imageID

	tx := db.Begin()

	if err := tx.Debug().Where("image_id = ?", imageID).FirstOrCreate(&i).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&i).Updates(map[string]interface{}{"json": json, "uploaded": false, "checksumed": false}).Error; err != nil {
		tx.Rollback()
		return err
	} else if err == nil {
		tx.Commit()
		return nil
	}

	tx.Commit()
	return nil
}

//PutLayer is put image layer, path, uploaded and size.
func (i *DockerImageV1) PutLayer(imageID, path string, size int64) error {
	tx := db.Begin()

	if err := tx.Debug().Where("image_id = ?", imageID).First(&i).Updates(map[string]interface{}{"path": path, "uploaded": true, "size": size}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}

//PutChecksum is put image's checksum, payload and ancestry.
func (i *DockerImageV1) PutChecksum(imageID, checksum, payload string) error {
	tx := db.Begin()

	var data map[string]interface{}
	var ancestries []string
	var parentAnestries []string

	if err := tx.Debug().Where("image_id = ?", imageID).First(&i).Error; err != nil {
		tx.Rollback()
		return err
	}

	ancestries = append(ancestries, imageID)

	if err := json.Unmarshal([]byte(i.JSON), &data); err != nil {
		tx.Rollback()
		return err
	}

	if value, has := data["parent"]; has == true {
		image := new(DockerImageV1)

		if err := tx.Debug().Where("image_id = ?", value.(string)).First(&image).Error; err != nil {
			tx.Rollback()
			return err
		}

		if err := json.Unmarshal([]byte(image.Ancestry), &parentAnestries); err != nil {
			tx.Rollback()
			return err
		}

		ancestries = append(ancestries, parentAnestries...)
	}

	ancestry, _ := json.Marshal(ancestries)

	if err := tx.Debug().Model(&i).Updates(map[string]interface{}{"checksum": checksum, "payload": payload, "ancestry": string(ancestry)}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}

//Put is set tag in the database.
func (t *DockerTagV1) Put(imageID, tag, namespace, repository string) error {
	tx := db.Begin()

	r := new(DockerV1)
	if err := tx.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		tx.Rollback()
		return err
	}

	t.DockerV1 = r.ID
	t.ImageID = imageID
	t.Tag = tag
	if err := tx.Debug().Where("docker_v1 = ? AND tag = ?", r.ID, tag).FirstOrCreate(&t).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&t).Updates(map[string]interface{}{"image_id": imageID}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}
<commit_msg>Add short description in the DockerV1 repository.<commit_after>/*
Copyright 2015 The ContainerOps Authors All rights reserved.

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 models

import (
	"encoding/json"
	"time"
)

//DockerV1 is Docker Repository V1 repository.
type DockerV1 struct {
	ID          int64      `json:"id" gorm:"primary_key"`
	Namespace   string     `json:"namespace" sql:"not null;type:varchar(255)" gorm:"unique_index:dockerv1_repository"`
	Repository  string     `json:"repository" sql:"not null;type:varchar(255)" gorm:"unique_index:dockerv1_repository"`
	JSON        string     `json:"json" sql:"null;type:text"`
	Manifests   string     `json:"manifests" sql:"null;type:text"`
	Agent       string     `json:"agent" sql:"null;type:text"`
  Short       string     `json:"short" sql:"null;type:text"`
	Description string     `json:"description" sql:"null;type:text"`
	Size        int64      `json:"size" sql:"default:0"`
	Locked      bool       `json:"locked" sql:"default:false"` //When create/update the repository, the locked will be true.
	CreatedAt   time.Time  `json:"create_at" sql:""`
	UpdatedAt   time.Time  `json:"update_at" sql:""`
	DeletedAt   *time.Time `json:"delete_at" sql:"index"`
}

//TableName in mysql is "docker_v1".
func (r *DockerV1) TableName() string {
	return "docker_v1"
}

//DockerImageV1 is
type DockerImageV1 struct {
	ID         int64      `json:"id" gorm:"primary_key"`
	ImageID    string     `json:"image_id" sql:"not null;unique;varchar(255)"`
	JSON       string     `json:"json" sql:"null;type:text"`
	Ancestry   string     `json:"ancestry" sql:"null;type:text"`
	Checksum   string     `json:"checksum" sql:"null;type:varchar(255)"`
	Payload    string     `json:"payload" sql:"null;type:varchar(255)"`
	Path       string     `json:"path" sql:"null;type:text"`
	OSS        string     `json:"oss" sql:"null;type:text"`
	Size       int64      `json:"size" sql:"default:0"`
	Uploaded   bool       `json:"uploaded" sql:"default:false"`
	Checksumed bool       `json:"checksumed" sql:"default:false"`
	Locked     bool       `json:"locked" sql:"default:false"`
	CreatedAt  time.Time  `json:"create_at" sql:""`
	UpdatedAt  time.Time  `json:"update_at" sql:""`
	DeletedAt  *time.Time `json:"delete_at" sql:"index"`
}

//TableName in mysql is "docker_image_v1".
func (i *DockerImageV1) TableName() string {
	return "docker_image_v1"
}

//DockerTagV1 is
type DockerTagV1 struct {
	ID        int64      `json:"id" gorm:"primary_key"`
	DockerV1  int64      `json:"docker_v1" sql:"not null;default:0"`
	Tag       string     `json:"tag" sql:"not null;varchar(255)"`
	ImageID   string     `json:"image_id" sql:"not null;varchar(255)"`
	CreatedAt time.Time  `json:"create_at" sql:""`
	UpdatedAt time.Time  `json:"update_at" sql:""`
	DeletedAt *time.Time `json:"delete_at" sql:"index"`
}

//TableName in mysql is "docker_tag_v1".
func (t *DockerTagV1) TableName() string {
	return "docker_tag_v1"
}

//Put function will create or update repository.
func (r *DockerV1) Put(namespace, repository, json, agent string) error {
	r.Namespace, r.Repository, r.JSON, r.Agent, r.Locked = namespace, repository, json, agent, true

	tx := db.Begin()

	if err := tx.Debug().Where("namespace = ? AND repository = ? ", namespace, repository).FirstOrCreate(&r).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&r).Updates(map[string]interface{}{"json": json, "agent": agent, "locked": true}).Error; err != nil {
		tx.Rollback()
		return err
	} else if err == nil {
		tx.Commit()
		return nil
	}

	tx.Commit()
	return nil
}

//Unlocked is Unlocked repository data so could pull.
func (r *DockerV1) Unlocked(namespace, repository string) error {
	tx := db.Begin()

	if err := tx.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&r).Updates(map[string]interface{}{"locked": false}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}

//Get return Docker V1 repository data.
func (r *DockerV1) Get(namespace, repository string) (DockerV1, error) {
	if err := db.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		return *new(DockerV1), err
	} else {
		return *r, nil
	}
}

//GetTags return tas data of repository.
func (r *DockerV1) GetTags(namespace, repository string) (map[string]string, error) {
	if err := db.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		return map[string]string{}, err
	} else {
		var tags []DockerTagV1
		result := map[string]string{}

		if err := db.Debug().Where("docker_v1 = ?", r.ID).Find(&tags).Error; err != nil {
			return map[string]string{}, err
		}

		for _, tag := range tags {
			result[tag.Tag] = tag.ImageID
		}

		return result, nil
	}
}

//Get is search image by ImageID.
func (i *DockerImageV1) Get(imageID string) (DockerImageV1, error) {
	if err := db.Debug().Where("image_id = ?", imageID).First(&i).Error; err != nil {
		return *i, err
	} else {
		return *i, nil
	}
}

//PutJSON is put image json by ImageID.
func (i *DockerImageV1) PutJSON(imageID, json string) error {
	i.ImageID = imageID

	tx := db.Begin()

	if err := tx.Debug().Where("image_id = ?", imageID).FirstOrCreate(&i).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&i).Updates(map[string]interface{}{"json": json, "uploaded": false, "checksumed": false}).Error; err != nil {
		tx.Rollback()
		return err
	} else if err == nil {
		tx.Commit()
		return nil
	}

	tx.Commit()
	return nil
}

//PutLayer is put image layer, path, uploaded and size.
func (i *DockerImageV1) PutLayer(imageID, path string, size int64) error {
	tx := db.Begin()

	if err := tx.Debug().Where("image_id = ?", imageID).First(&i).Updates(map[string]interface{}{"path": path, "uploaded": true, "size": size}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}

//PutChecksum is put image's checksum, payload and ancestry.
func (i *DockerImageV1) PutChecksum(imageID, checksum, payload string) error {
	tx := db.Begin()

	var data map[string]interface{}
	var ancestries []string
	var parentAnestries []string

	if err := tx.Debug().Where("image_id = ?", imageID).First(&i).Error; err != nil {
		tx.Rollback()
		return err
	}

	ancestries = append(ancestries, imageID)

	if err := json.Unmarshal([]byte(i.JSON), &data); err != nil {
		tx.Rollback()
		return err
	}

	if value, has := data["parent"]; has == true {
		image := new(DockerImageV1)

		if err := tx.Debug().Where("image_id = ?", value.(string)).First(&image).Error; err != nil {
			tx.Rollback()
			return err
		}

		if err := json.Unmarshal([]byte(image.Ancestry), &parentAnestries); err != nil {
			tx.Rollback()
			return err
		}

		ancestries = append(ancestries, parentAnestries...)
	}

	ancestry, _ := json.Marshal(ancestries)

	if err := tx.Debug().Model(&i).Updates(map[string]interface{}{"checksum": checksum, "payload": payload, "ancestry": string(ancestry)}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}

//Put is set tag in the database.
func (t *DockerTagV1) Put(imageID, tag, namespace, repository string) error {
	tx := db.Begin()

	r := new(DockerV1)
	if err := tx.Debug().Where("namespace = ? AND repository = ?", namespace, repository).First(&r).Error; err != nil {
		tx.Rollback()
		return err
	}

	t.DockerV1 = r.ID
	t.ImageID = imageID
	t.Tag = tag
	if err := tx.Debug().Where("docker_v1 = ? AND tag = ?", r.ID, tag).FirstOrCreate(&t).Error; err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Debug().Model(&t).Updates(map[string]interface{}{"image_id": imageID}).Error; err != nil {
		tx.Rollback()
		return err
	}

	tx.Commit()
	return nil
}
<|endoftext|>


<commit_before>package main

import "fmt"

const (
	retryOrContactSupport        = "Please try again, or contact support@koding.com"
	retryNewCodeOrContactSupport = `Please go back to Koding to get a new code and try again, or contact support@koding.com`
)

var (
	// InternalRetryError is typically used when the circumstances that caused the error
	// are implementation details only and cannot be told to the user, and that they
	// can't do anything to correct the error.
	//
	// The only thing left is to .. retry, or possibly report to support@koding.com?
	GenericInternalError = fmt.Sprintf(
		"Error: Encountered an internal error.\n%s", retryOrContactSupport,
	)

	GenericInternalErrorRetry = fmt.Sprintf(
		"Error: Encountered an internal error.\n%s", retryNewCodeOrContactSupport,
	)

	FailedInstallingKlient = fmt.Sprintf(
		"Error: Unable to install the %s.\n%s", KlientName, retryNewCodeOrContactSupport,
	)

	FailedDownloadingKlient = fmt.Sprintf(
		"Error: Unable to download the %s binary.\n%s",
		KlientName, retryNewCodeOrContactSupport,
	)

	FailedRegisteringKlient = fmt.Sprintf(
		"Error: Unable to authenticate %s to koding.com.\n%s",
		Name, retryNewCodeOrContactSupport,
	)

	FailedVerifyingInstall = fmt.Sprintf(
		"Error: Unable to verify the installation of %s.\n%s",
		Name, retryNewCodeOrContactSupport,
	)

	FailedStartKlient = fmt.Sprintf(
		"Error: Failed to start the %s within the expected time.\n%s", KlientName, retryOrContactSupport,
	)

	FailedStopKlient = fmt.Sprintf(
		"Error: Failed to stop the %s within the expected time.\n%s", KlientName, retryOrContactSupport,
	)
)
<commit_msg>styleguide: Added comments for golint<commit_after>package main

import "fmt"

const (
	retryOrContactSupport        = "Please try again, or contact support@koding.com"
	retryNewCodeOrContactSupport = `Please go back to Koding to get a new code and try again, or contact support@koding.com`
)

var (
	// GenericInternalError is a generic error message. Typically used when we don't
	// want to reveal what exactly went wrong, like confusing implementation details.
	GenericInternalError = fmt.Sprintf(
		"Error: Encountered an internal error.\n%s", retryOrContactSupport,
	)

	// GenericInternalErrorRetry is a generic error message. Typically used when we
	// don't want to reveal what exactly went wrong, like confusing implementation
	// details.
	//
	// It instructs them to get a new code and try again.
	GenericInternalErrorRetry = fmt.Sprintf(
		"Error: Encountered an internal error.\n%s", retryNewCodeOrContactSupport,
	)

	// FailedInstallingKlient is generic for when a klient install fails.
	FailedInstallingKlient = fmt.Sprintf(
		"Error: Unable to install the %s.\n%s", KlientName, retryNewCodeOrContactSupport,
	)

	// FailedDownloadingKlient is used when downloading klient fails.
	FailedDownloadingKlient = fmt.Sprintf(
		"Error: Unable to download the %s binary.\n%s",
		KlientName, retryNewCodeOrContactSupport,
	)

	// FailedRegisteringKlient is used when registering klient to kontrol fails.
	FailedRegisteringKlient = fmt.Sprintf(
		"Error: Unable to authenticate %s to koding.com.\n%s",
		Name, retryNewCodeOrContactSupport,
	)

	// FailedVerifyingInstall is used when verifying the install fails.
	FailedVerifyingInstall = fmt.Sprintf(
		"Error: Unable to verify the installation of %s.\n%s",
		Name, retryNewCodeOrContactSupport,
	)

	// FailedStartKlient is used when starting klient fails.
	FailedStartKlient = fmt.Sprintf(
		"Error: Failed to start the %s within the expected time.\n%s", KlientName, retryOrContactSupport,
	)

	// FailedStopKlient is used when stopping klient fails.
	FailedStopKlient = fmt.Sprintf(
		"Error: Failed to stop the %s within the expected time.\n%s", KlientName, retryOrContactSupport,
	)
)
<|endoftext|>


<commit_before>// Package errors extends the errors package in the stdlib.
package errors

import "errors"

// New returns an error that formats as the given text.
//
// This is an alias to the stdlib errors.New function.
func New(text string) error {
	return errors.New(text)
}
<commit_msg>Update errors package documentation<commit_after>// Package errors extends the errors package in the stdlib.
//
// Despite the implicit nature of interface satisfication in Go
// this package exports a number of interfaces to avoid defining them over and over again.
// Although it means coupling between the consumer code and this package,
// the purpose of this library (being a stdlib extension) justifies that.
package errors

import "errors"

// New returns an error that formats as the given text.
//
// This is an alias to the stdlib errors.New function.
func New(text string) error {
	return errors.New(text)
}
<|endoftext|>


<commit_before>package models

import (
	"io/ioutil"

	"gopkg.in/yaml.v2"
)

type Router struct {
	Path        string
	AllowedHost []string
}

type Setting struct {
	Database struct {
		Host       string
		DbName     string
		TokenTable string
		UserTable  string
	}
	Ssl struct {
		Key         string
		Certificate string
	}
	Router struct {
		Register Router
		Login    Router
		Validate Router
		Logout   Router
	}
	Domain string
	IP     string
}

// Set the set of loaded settings.
var Set Setting

// LoadSettings loads the settings from a file.
func LoadSettings(path string) error {
	text, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	if err = yaml.Unmarshal(text, &Set); err != nil {
		return err
	}

	return nil
}
<commit_msg>Add update settings section.<commit_after>package models

import (
	"io/ioutil"

	"gopkg.in/yaml.v2"
)

type Router struct {
	Path        string
	AllowedHost []string
}

type Setting struct {
	Database struct {
		Host       string
		DbName     string
		TokenTable string
		UserTable  string
	}
	Ssl struct {
		Key         string
		Certificate string
	}
	Router struct {
		Register Router
		Login    Router
		Validate Router
		Logout   Router
		Update   Router
	}
	Domain string
	IP     string
}

// Set the set of loaded settings.
var Set Setting

// LoadSettings loads the settings from a file.
func LoadSettings(path string) error {
	text, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	if err = yaml.Unmarshal(text, &Set); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"errors"
	"strings"
	"strconv"
	"math"
	"os"
	"path"
	"io"
	"bufio"
	"container/list"
	"github.com/jessevdk/go-flags"
)

const (
	version = "0.1"
)

type subtitle struct {
	text string
	start uint
	end uint
}

func die(err error) {
	fmt.Fprintf(os.Stderr, "%v\n", err)
	os.Exit(1)
}

func roundFloat64(f float64) float64 {
	val := f - float64(int64(f))
	if val >= 0.5 {
		return math.Ceil(f)
	} else if val > 0 {
		return math.Floor(f)
	} else if val <= -0.5 {
		return math.Floor(f)
	} else if val < 0 {
		return math.Ceil(f)
	}
	return f
}

/* converts hh:mm:ss,mss to milliseconds */
func time_to_msecs(tm string) (uint, error) {
	var msecs uint
	var h, m, s, ms uint

	tm = strings.Replace(tm, ".", ",", 1)
	num, err := fmt.Sscanf(tm, "%d:%d:%d,%d", &h, &m, &s, &ms)

	if num != 4 || err != nil {
		return 0, errors.New("Parsing error: Can not covert `" + tm + "' to milliseconds.")
	}

	msecs = h * 60 * 60 * 1000
	msecs += m * 60 * 1000
	msecs += s * 1000
	msecs += ms

	return msecs, nil
}

/* converts milliseconds to hh:mm:ss,mss */
func msecs_to_time(msecs uint) string {
	var h, m, s, ms uint

	h = msecs / (60 * 60 * 1000)
	msecs %= 60 * 60 * 1000
	m = msecs / (60 * 1000)
	msecs %= 60 * 1000
	s = msecs / 1000
	ms = msecs % 1000

	tm := fmt.Sprintf("%02d:%02d:%02d,%03d", h, m, s, ms)

	return tm
}

/* read SubRip (srt) file */
func read_srt(filename string) (*list.List, error) {
	var state int = 0
	var subs *list.List
	var sub *subtitle

	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	r := bufio.NewReader(f)
	subs = list.New()
	sub = new(subtitle)

	for {
		var (
			isprefix bool = true
			err error = nil
			ln, line []byte
		)

		for isprefix && err == nil {
			line, isprefix, err = r.ReadLine()
			if err != nil && err != io.EOF {
				return nil, err
			}
			ln = append(ln, line...)
		}

		/* parse subtitle id */
		if state == 0 {
			/* avoid false-positive parsing error */
			if err == io.EOF && len(ln) == 0 {
				break;
			}
			id := strings.Split(string(ln), " ")
			if len(id) != 1 {
				return nil, errors.New("Parsing error: Wrong file format")
			}
			_, err = strconv.ParseUint(id[0], 10, 0)
			if err != nil {
				return nil, err
				return nil, errors.New("Parsing error: Wrong file format")
			}
			state = 1
		/* parse start, end times */
		} else if state == 1 {
			tm := strings.Split(string(ln), " ")
			if len(tm) != 3 || tm[1] != "-->" {
				return nil, errors.New("Parsing error: Wrong file format")
			}
			sub.start, err = time_to_msecs(tm[0])
			if err != nil {
				return nil, err
			}
			sub.end, err = time_to_msecs(tm[2])
			if err != nil {
				return nil, err
			}
			state = 2
		/* parse the actual subtitle text */
		} else if state == 2 {
			if len(ln) == 0 {
				subs.PushBack(sub)
				sub = new(subtitle)
				state = 0
			} else {
				sub.text += string(ln) + "\r\n"
			}
		}

		if err == io.EOF {
			break;
		}
	}

	return subs, nil
}

/* write SubRip (srt) file */
func write_srt(filename string, subs *list.List) error {
	var id int = 0

	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	w := bufio.NewWriter(f)
	defer w.Flush()

	for e := subs.Front(); e != nil; e = e.Next() {
		id++
		sub := e.Value.(*subtitle)
		fmt.Fprintf(w, "%d\r\n", id)
		fmt.Fprintf(w, "%s --> %s\r\n", msecs_to_time(sub.start), msecs_to_time(sub.end))
		fmt.Fprintf(w, "%s\r\n", sub.text)
	}

	return nil
}

/* synchronize subtitles by knowing the time of the first and the last subtitle.
 * to archive this we must use the linear equation: y = mx + b */
func sync_subs(subs *list.List, synced_first_ms uint, synced_last_ms uint) {
	var slope, yint float64

	desynced_first_ms := subs.Front().Value.(*subtitle).start
	desynced_last_ms := subs.Back().Value.(*subtitle).start

	/* m = (y2 - y1) / (x2 - x1)
	 * m: slope
	 * y2: synced_last_ms
	 * y1: synced_first_ms
	 * x2: desynced_last_ms
	 * x1: desynced_first_ms */
	slope = float64(synced_last_ms - synced_first_ms) / float64(desynced_last_ms - desynced_first_ms)
	/* b = y - mx
	 * b: yint
	 * y: synced_last_ms
	 * m: slope
	 * x: desynced_last_ms */
	yint = float64(synced_last_ms) - slope * float64(desynced_last_ms)

	for e := subs.Front(); e != nil; e = e.Next() {
		sub := e.Value.(*subtitle)
		/* y = mx + b
		 * y: sub.start and sub.end
		 * m: slope
		 * x: sub.start and sub.end
		 * b: yint */
		sub.start = uint(roundFloat64(slope * float64(sub.start) + yint))
		sub.end = uint(roundFloat64(slope * float64(sub.end) + yint))
	}
}

func main() {
	var first_ms, last_ms uint

	var opts struct {
		FirstTm string `short:"f" long:"first-sub" description:"Time of first subtitle"`
		LastTm string `short:"l" long:"last-sub" description:"Time of last subtitle"`
		InputFl string `short:"i" long:"input" description:"Input file"`
		OutputFl string `short:"o" long:"output" description:"Output file"`
		PrintVersion bool `short:"v" long:"version" description:"Print version"`
	}

	_, err := flags.Parse(&opts)
	if err != nil {
		if err.(*flags.Error).Type == flags.ErrHelp {
			fmt.Fprintf(os.Stderr, "Example:\n")
			fmt.Fprintf(os.Stderr, "  %s -f 00:01:33,492 -l 01:39:23,561 -i file.srt\n",
				path.Base(os.Args[0]))
			os.Exit(0)
		}
		os.Exit(1)
	}

	if opts.PrintVersion {
		fmt.Printf("subsync v%s\n", version)
		os.Exit(0)
	}

	if opts.InputFl == "" {
		fmt.Fprintf(os.Stderr, "You must specify an input file with -i option.\n")
		os.Exit(1)
	}

	if opts.FirstTm != "" {
		first_ms, err = time_to_msecs(opts.FirstTm)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Please check the value of -f option.\n")
			die(err)
		}
	}

	if opts.LastTm != "" {
		last_ms, err = time_to_msecs(opts.LastTm)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Please check the value of -l option.\n")
			die(err)
		}
	}

	/* if output file is not set, use the input file */
	if opts.OutputFl == "" {
		opts.OutputFl = opts.InputFl
	}

	subs, err := read_srt(opts.InputFl)
	if err != nil {
		die(err)
	}

	/* if time of the first synced subtitle is not set,
	 * use the time of the first desynced subtitle */
	if opts.FirstTm == "" {
		first_ms = subs.Front().Value.(*subtitle).start
	}

	/* if time of the last synced subtitle is not set,
	 * use the time of the last desynced subtitle */
	if opts.LastTm == "" {
		last_ms = subs.Back().Value.(*subtitle).start
	}

	if first_ms > last_ms {
		fmt.Fprintf(os.Stderr, "First subtitle can not be after last subtitle.\n")
		fmt.Fprintf(os.Stderr, "Please check the values of -f and/or -l options.\n")
		os.Exit(1)
	}

	sync_subs(subs, first_ms, last_ms)

	err = write_srt(opts.OutputFl, subs)
	if err != nil {
		die(err)
	}
}<commit_msg>remove a forgotten line<commit_after>package main

import (
	"fmt"
	"errors"
	"strings"
	"strconv"
	"math"
	"os"
	"path"
	"io"
	"bufio"
	"container/list"
	"github.com/jessevdk/go-flags"
)

const (
	version = "0.1"
)

type subtitle struct {
	text string
	start uint
	end uint
}

func die(err error) {
	fmt.Fprintf(os.Stderr, "%v\n", err)
	os.Exit(1)
}

func roundFloat64(f float64) float64 {
	val := f - float64(int64(f))
	if val >= 0.5 {
		return math.Ceil(f)
	} else if val > 0 {
		return math.Floor(f)
	} else if val <= -0.5 {
		return math.Floor(f)
	} else if val < 0 {
		return math.Ceil(f)
	}
	return f
}

/* converts hh:mm:ss,mss to milliseconds */
func time_to_msecs(tm string) (uint, error) {
	var msecs uint
	var h, m, s, ms uint

	tm = strings.Replace(tm, ".", ",", 1)
	num, err := fmt.Sscanf(tm, "%d:%d:%d,%d", &h, &m, &s, &ms)

	if num != 4 || err != nil {
		return 0, errors.New("Parsing error: Can not covert `" + tm + "' to milliseconds.")
	}

	msecs = h * 60 * 60 * 1000
	msecs += m * 60 * 1000
	msecs += s * 1000
	msecs += ms

	return msecs, nil
}

/* converts milliseconds to hh:mm:ss,mss */
func msecs_to_time(msecs uint) string {
	var h, m, s, ms uint

	h = msecs / (60 * 60 * 1000)
	msecs %= 60 * 60 * 1000
	m = msecs / (60 * 1000)
	msecs %= 60 * 1000
	s = msecs / 1000
	ms = msecs % 1000

	tm := fmt.Sprintf("%02d:%02d:%02d,%03d", h, m, s, ms)

	return tm
}

/* read SubRip (srt) file */
func read_srt(filename string) (*list.List, error) {
	var state int = 0
	var subs *list.List
	var sub *subtitle

	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	r := bufio.NewReader(f)
	subs = list.New()
	sub = new(subtitle)

	for {
		var (
			isprefix bool = true
			err error = nil
			ln, line []byte
		)

		for isprefix && err == nil {
			line, isprefix, err = r.ReadLine()
			if err != nil && err != io.EOF {
				return nil, err
			}
			ln = append(ln, line...)
		}

		/* parse subtitle id */
		if state == 0 {
			/* avoid false-positive parsing error */
			if err == io.EOF && len(ln) == 0 {
				break;
			}
			id := strings.Split(string(ln), " ")
			if len(id) != 1 {
				return nil, errors.New("Parsing error: Wrong file format")
			}
			_, err = strconv.ParseUint(id[0], 10, 0)
			if err != nil {
				return nil, errors.New("Parsing error: Wrong file format")
			}
			state = 1
		/* parse start, end times */
		} else if state == 1 {
			tm := strings.Split(string(ln), " ")
			if len(tm) != 3 || tm[1] != "-->" {
				return nil, errors.New("Parsing error: Wrong file format")
			}
			sub.start, err = time_to_msecs(tm[0])
			if err != nil {
				return nil, err
			}
			sub.end, err = time_to_msecs(tm[2])
			if err != nil {
				return nil, err
			}
			state = 2
		/* parse the actual subtitle text */
		} else if state == 2 {
			if len(ln) == 0 {
				subs.PushBack(sub)
				sub = new(subtitle)
				state = 0
			} else {
				sub.text += string(ln) + "\r\n"
			}
		}

		if err == io.EOF {
			break;
		}
	}

	return subs, nil
}

/* write SubRip (srt) file */
func write_srt(filename string, subs *list.List) error {
	var id int = 0

	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	w := bufio.NewWriter(f)
	defer w.Flush()

	for e := subs.Front(); e != nil; e = e.Next() {
		id++
		sub := e.Value.(*subtitle)
		fmt.Fprintf(w, "%d\r\n", id)
		fmt.Fprintf(w, "%s --> %s\r\n", msecs_to_time(sub.start), msecs_to_time(sub.end))
		fmt.Fprintf(w, "%s\r\n", sub.text)
	}

	return nil
}

/* synchronize subtitles by knowing the time of the first and the last subtitle.
 * to archive this we must use the linear equation: y = mx + b */
func sync_subs(subs *list.List, synced_first_ms uint, synced_last_ms uint) {
	var slope, yint float64

	desynced_first_ms := subs.Front().Value.(*subtitle).start
	desynced_last_ms := subs.Back().Value.(*subtitle).start

	/* m = (y2 - y1) / (x2 - x1)
	 * m: slope
	 * y2: synced_last_ms
	 * y1: synced_first_ms
	 * x2: desynced_last_ms
	 * x1: desynced_first_ms */
	slope = float64(synced_last_ms - synced_first_ms) / float64(desynced_last_ms - desynced_first_ms)
	/* b = y - mx
	 * b: yint
	 * y: synced_last_ms
	 * m: slope
	 * x: desynced_last_ms */
	yint = float64(synced_last_ms) - slope * float64(desynced_last_ms)

	for e := subs.Front(); e != nil; e = e.Next() {
		sub := e.Value.(*subtitle)
		/* y = mx + b
		 * y: sub.start and sub.end
		 * m: slope
		 * x: sub.start and sub.end
		 * b: yint */
		sub.start = uint(roundFloat64(slope * float64(sub.start) + yint))
		sub.end = uint(roundFloat64(slope * float64(sub.end) + yint))
	}
}

func main() {
	var first_ms, last_ms uint

	var opts struct {
		FirstTm string `short:"f" long:"first-sub" description:"Time of first subtitle"`
		LastTm string `short:"l" long:"last-sub" description:"Time of last subtitle"`
		InputFl string `short:"i" long:"input" description:"Input file"`
		OutputFl string `short:"o" long:"output" description:"Output file"`
		PrintVersion bool `short:"v" long:"version" description:"Print version"`
	}

	_, err := flags.Parse(&opts)
	if err != nil {
		if err.(*flags.Error).Type == flags.ErrHelp {
			fmt.Fprintf(os.Stderr, "Example:\n")
			fmt.Fprintf(os.Stderr, "  %s -f 00:01:33,492 -l 01:39:23,561 -i file.srt\n",
				path.Base(os.Args[0]))
			os.Exit(0)
		}
		os.Exit(1)
	}

	if opts.PrintVersion {
		fmt.Printf("subsync v%s\n", version)
		os.Exit(0)
	}

	if opts.InputFl == "" {
		fmt.Fprintf(os.Stderr, "You must specify an input file with -i option.\n")
		os.Exit(1)
	}

	if opts.FirstTm != "" {
		first_ms, err = time_to_msecs(opts.FirstTm)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Please check the value of -f option.\n")
			die(err)
		}
	}

	if opts.LastTm != "" {
		last_ms, err = time_to_msecs(opts.LastTm)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Please check the value of -l option.\n")
			die(err)
		}
	}

	/* if output file is not set, use the input file */
	if opts.OutputFl == "" {
		opts.OutputFl = opts.InputFl
	}

	subs, err := read_srt(opts.InputFl)
	if err != nil {
		die(err)
	}

	/* if time of the first synced subtitle is not set,
	 * use the time of the first desynced subtitle */
	if opts.FirstTm == "" {
		first_ms = subs.Front().Value.(*subtitle).start
	}

	/* if time of the last synced subtitle is not set,
	 * use the time of the last desynced subtitle */
	if opts.LastTm == "" {
		last_ms = subs.Back().Value.(*subtitle).start
	}

	if first_ms > last_ms {
		fmt.Fprintf(os.Stderr, "First subtitle can not be after last subtitle.\n")
		fmt.Fprintf(os.Stderr, "Please check the values of -f and/or -l options.\n")
		os.Exit(1)
	}

	sync_subs(subs, first_ms, last_ms)

	err = write_srt(opts.OutputFl, subs)
	if err != nil {
		die(err)
	}
}<|endoftext|>


<commit_before>// Package svc provides some tooling to make building services with remind101/pkg
// easier.
//
// Recommend Usage:
//
//	func main() {
//		env := svc.InitAll()
//		defer env.Close()
//
//		r := httpx.NewRouter()
//		// ... add routes
//
//		h := svc.NewStandardHandler(svc.HandlerOpts{
//			Router:   r,
//			Reporter: env.Reporter,
//	})
//
// 	svc.RunServer(h, "80", 5*time.Second)
// }
package svc

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	ddtrace "github.com/DataDog/dd-trace-go/opentracing"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/pkg/errors"
	"github.com/remind101/pkg/httpx"
	"github.com/remind101/pkg/httpx/middleware"
	"github.com/remind101/pkg/logger"
	"github.com/remind101/pkg/metrics"
	"github.com/remind101/pkg/reporter"
	"github.com/remind101/pkg/reporter/rollbar"
)

type HandlerOpts struct {
	Router            *httpx.Router
	Reporter          reporter.Reporter
	ForwardingHeaders []string
	BasicAuth         string
	ErrorHandler      middleware.ErrorHandlerFunc
	HandlerTimeout    time.Duration
}

// NewStandardHandler returns an http.Handler with a standard middleware stack.
// The last middleware added is the first middleware to handle the request.
// Order is pretty important as some middleware depends on others having run
// already.
func NewStandardHandler(opts HandlerOpts) http.Handler {
	h := httpx.Handler(opts.Router)

	if opts.HandlerTimeout != 0 {
		// Timeout requests after the given Timeout duration.
		h = middleware.TimeoutHandler(h, opts.HandlerTimeout)
	}

	// Recover from panics. A panic is converted to an error. This should be first,
	// even though it means panics in middleware will not be recovered, because
	// later middleware expects endpoint panics to be returned as an error.
	h = middleware.BasicRecover(h)

	// Handler errors returned by endpoint handler or recovery middleware.
	// Errors will no longer be returned after this middeware.
	errorHandler := opts.ErrorHandler
	if errorHandler == nil {
		errorHandler = middleware.ReportingErrorHandler
	}
	h = middleware.HandleError(h, errorHandler)

	// Add request tracing. Must go after the HandleError middleware in order
	// to capture the status code written to the response.
	h = middleware.OpentracingTracing(h, opts.Router)

	// Insert logger into context and log requests at INFO level.
	h = middleware.LogTo(h, middleware.LoggerWithRequestID)

	// Add reporter to context and request to reporter context.
	h = middleware.WithReporter(h, opts.Reporter)

	// Add the request id to the context.
	h = middleware.ExtractRequestID(h)

	// Add basic auth
	if opts.BasicAuth != "" {
		user := strings.Split(opts.BasicAuth, ":")[0]
		pass := strings.Split(opts.BasicAuth, ":")[1]
		h = middleware.BasicAuth(h, user, pass, "")
	}

	// Adds forwarding headers from request to the context. This allows http clients
	// to get those headers from the context and add them to upstream requests.
	if len(opts.ForwardingHeaders) > 0 {
		for _, header := range opts.ForwardingHeaders {
			h = middleware.ExtractHeader(h, header)
		}
	}

	// Wrap the route in middleware to add a context.Context. This middleware must be
	// last as it acts as the adaptor between http.Handler and httpx.Handler.
	return middleware.BackgroundContext(h)
}

// RunServer handles the biolerplate of starting an http server and handling
// signals gracefully.
func RunServer(h http.Handler, port string, writeTimeout time.Duration) {
	errCh := make(chan error)

	// Handle SIGINT and SIGTERM.
	sigCh := make(chan os.Signal)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	fmt.Printf("Listening on port %s\n", port)

	// Add timeouts to the server
	srv := &http.Server{
		WriteTimeout: writeTimeout * time.Second,
		Addr:         ":" + port,
		Handler:      h,
	}

	go func() {
		defer reporter.Monitor(context.Background())
		err := srv.ListenAndServe()
		if err != nil {
			errCh <- errors.Wrapf(err, "unable to start server")
		}
	}()

	select {
	case sig := <-sigCh:
		fmt.Println("Received signal, stopping.", "signal", sig)
	// Cleanup
	case err := <-errCh:
		fmt.Println(err)
		os.Exit(1)
	}
}

// Env holds global dependencies that need to be initialized in main() and
// injected as dependencies into an application.
type Env struct {
	Reporter reporter.Reporter
	Logger   logger.Logger
	Context  context.Context
	Close    func() // Should be called in a defer in main().
}

// InitAll will initialize all the common dependencies such as metrics, reporting,
// tracing, and logging.
func InitAll() Env {
	traceCloser := InitTracer()
	metricsCloser := InitMetrics()

	l := InitLogger()
	logger.DefaultLogger = l

	r := InitReporter()

	ctx := reporter.WithReporter(context.Background(), r)
	ctx = logger.WithLogger(ctx, l)

	go func() {
		defer reporter.Monitor(ctx)
		metrics.Runtime()
	}()

	return Env{
		Logger:   l,
		Reporter: r,
		Context:  ctx,
		Close: func() {
			traceCloser()
			metricsCloser()
			reporter.Monitor(ctx)
		},
	}
}

// InitTracer configures a global datadog tracer.
//
// Env Vars:
// * DDTRACE_ADDR - The host:port of the local trace agent server.
// * EMPIRE_APPNAME - App name, used to construct the service name.
// * EMPIRE_PROCESS - Process name, used to construct the service name.
func InitTracer() func() {
	// create a Tracer configuration
	config := ddtrace.NewConfiguration()
	config.ServiceName = fmt.Sprintf("%s.%s", os.Getenv("EMPIRE_APPNAME"), os.Getenv("EMPIRE_PROCESS"))
	if addr := os.Getenv("DDTRACE_ADDR"); addr != "" {
		config.AgentHostname = addr
	}

	// Initialize a Tracer and ensure a graceful shutdown
	// using the `closer.Close()`
	tracer, closer, err := ddtrace.NewTracer(config)
	if err != nil {
		fmt.Println(err)
	}

	// set the Datadog tracer as a GlobalTracer
	opentracing.SetGlobalTracer(tracer)

	return func() {
		closer.Close()
	}
}

// InitMetrics configures pkg/metrics
//
// Env Vars:
// * STATSD_ADDR - The host:port of the statsd server.
func InitMetrics() func() {
	if addr := os.Getenv("STATSD_ADDR"); addr != "" {
		metrics.SetEmpireDefaultTags()
		metrics.Reporter, _ = metrics.NewDataDogMetricsReporter(addr)
	}

	return func() {
		metrics.Close()
	}
}

// InitLogger configures a leveled logger.
//
// Env Vars:
// * LOG_LEVEL - The log level
//
// If you want to replace the global default logger:
//	logger.DefaultLogger = InitLogger()
func InitLogger() logger.Logger {
	lvl := logger.ERROR
	if ll := os.Getenv("LOG_LEVEL"); ll != "" {
		lvl = logger.ParseLevel(ll)
	}

	return logger.New(log.New(os.Stdout, "", 0), lvl)
}

// InitReporter configures and returns a reporter.Reporter instance.
//
// Env Vars:
// * ROLLBAR_ACCESS_TOKEN - The Rollbar access token
// * ROLLBAR_ENVIRONMENT  - The Rollbar environment (staging, production)
func InitReporter() reporter.Reporter {
	rbToken := os.Getenv("ROLLBAR_ACCESS_TOKEN")
	rbEnv := os.Getenv("ROLLBAR_ENVIRONMENT")

	rep := reporter.MultiReporter{}

	// Log Reporter, uses package level logger.
	rep = append(rep, reporter.NewLogReporter())

	// Rollbar reporter
	if rbToken != "" && rbEnv != "" {
		rollbar.ConfigureReporter(rbToken, rbEnv)
		rep = append(rep, rollbar.Reporter)
	} else {
		fmt.Println("Rollbar is not configured, skipping Rollbar reporter")
	}

	return rep
}
<commit_msg>Add graceful shutdown<commit_after>// Package svc provides some tooling to make building services with remind101/pkg
// easier.
//
// Recommend Usage:
//
//	func main() {
//		env := svc.InitAll()
//		defer env.Close()
//
//		r := httpx.NewRouter()
//		// ... add routes
//
//		h := svc.NewStandardHandler(svc.HandlerOpts{
//			Router:   r,
//			Reporter: env.Reporter,
//	})
//
// 	svc.RunServer(h, "80", 5*time.Second)
// }
package svc

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	ddtrace "github.com/DataDog/dd-trace-go/opentracing"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/pkg/errors"
	"github.com/remind101/pkg/httpx"
	"github.com/remind101/pkg/httpx/middleware"
	"github.com/remind101/pkg/logger"
	"github.com/remind101/pkg/metrics"
	"github.com/remind101/pkg/reporter"
	"github.com/remind101/pkg/reporter/rollbar"
)

type HandlerOpts struct {
	Router            *httpx.Router
	Reporter          reporter.Reporter
	ForwardingHeaders []string
	BasicAuth         string
	ErrorHandler      middleware.ErrorHandlerFunc
	HandlerTimeout    time.Duration
}

// NewStandardHandler returns an http.Handler with a standard middleware stack.
// The last middleware added is the first middleware to handle the request.
// Order is pretty important as some middleware depends on others having run
// already.
func NewStandardHandler(opts HandlerOpts) http.Handler {
	h := httpx.Handler(opts.Router)

	if opts.HandlerTimeout != 0 {
		// Timeout requests after the given Timeout duration.
		h = middleware.TimeoutHandler(h, opts.HandlerTimeout)
	}

	// Recover from panics. A panic is converted to an error. This should be first,
	// even though it means panics in middleware will not be recovered, because
	// later middleware expects endpoint panics to be returned as an error.
	h = middleware.BasicRecover(h)

	// Handler errors returned by endpoint handler or recovery middleware.
	// Errors will no longer be returned after this middeware.
	errorHandler := opts.ErrorHandler
	if errorHandler == nil {
		errorHandler = middleware.ReportingErrorHandler
	}
	h = middleware.HandleError(h, errorHandler)

	// Add request tracing. Must go after the HandleError middleware in order
	// to capture the status code written to the response.
	h = middleware.OpentracingTracing(h, opts.Router)

	// Insert logger into context and log requests at INFO level.
	h = middleware.LogTo(h, middleware.LoggerWithRequestID)

	// Add reporter to context and request to reporter context.
	h = middleware.WithReporter(h, opts.Reporter)

	// Add the request id to the context.
	h = middleware.ExtractRequestID(h)

	// Add basic auth
	if opts.BasicAuth != "" {
		user := strings.Split(opts.BasicAuth, ":")[0]
		pass := strings.Split(opts.BasicAuth, ":")[1]
		h = middleware.BasicAuth(h, user, pass, "")
	}

	// Adds forwarding headers from request to the context. This allows http clients
	// to get those headers from the context and add them to upstream requests.
	if len(opts.ForwardingHeaders) > 0 {
		for _, header := range opts.ForwardingHeaders {
			h = middleware.ExtractHeader(h, header)
		}
	}

	// Wrap the route in middleware to add a context.Context. This middleware must be
	// last as it acts as the adaptor between http.Handler and httpx.Handler.
	return middleware.BackgroundContext(h)
}

// RunServer handles the biolerplate of starting an http server and handling
// signals gracefully.
func RunServer(h http.Handler, port string, writeTimeout time.Duration) {
	errCh := make(chan error)

	// Handle SIGINT and SIGTERM.
	sigCh := make(chan os.Signal)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	fmt.Printf("Listening on port %s\n", port)

	// Add timeouts to the server
	srv := &http.Server{
		WriteTimeout: writeTimeout * time.Second,
		Addr:         ":" + port,
		Handler:      h,
	}

	go func() {
		defer reporter.Monitor(context.Background())
		err := srv.ListenAndServe()
		if err != nil {
			errCh <- errors.Wrapf(err, "unable to start server")
		}
	}()

	select {
	case sig := <-sigCh:
		fmt.Println("Received signal, stopping.", "signal", sig)
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		srv.Shutdown(ctx)
	// Cleanup
	case err := <-errCh:
		fmt.Println(err)
		os.Exit(1)
	}
}

// Env holds global dependencies that need to be initialized in main() and
// injected as dependencies into an application.
type Env struct {
	Reporter reporter.Reporter
	Logger   logger.Logger
	Context  context.Context
	Close    func() // Should be called in a defer in main().
}

// InitAll will initialize all the common dependencies such as metrics, reporting,
// tracing, and logging.
func InitAll() Env {
	traceCloser := InitTracer()
	metricsCloser := InitMetrics()

	l := InitLogger()
	logger.DefaultLogger = l

	r := InitReporter()

	ctx := reporter.WithReporter(context.Background(), r)
	ctx = logger.WithLogger(ctx, l)

	go func() {
		defer reporter.Monitor(ctx)
		metrics.Runtime()
	}()

	return Env{
		Logger:   l,
		Reporter: r,
		Context:  ctx,
		Close: func() {
			traceCloser()
			metricsCloser()
			reporter.Monitor(ctx)
		},
	}
}

// InitTracer configures a global datadog tracer.
//
// Env Vars:
// * DDTRACE_ADDR - The host:port of the local trace agent server.
// * EMPIRE_APPNAME - App name, used to construct the service name.
// * EMPIRE_PROCESS - Process name, used to construct the service name.
func InitTracer() func() {
	// create a Tracer configuration
	config := ddtrace.NewConfiguration()
	config.ServiceName = fmt.Sprintf("%s.%s", os.Getenv("EMPIRE_APPNAME"), os.Getenv("EMPIRE_PROCESS"))
	if addr := os.Getenv("DDTRACE_ADDR"); addr != "" {
		config.AgentHostname = addr
	}

	// Initialize a Tracer and ensure a graceful shutdown
	// using the `closer.Close()`
	tracer, closer, err := ddtrace.NewTracer(config)
	if err != nil {
		fmt.Println(err)
	}

	// set the Datadog tracer as a GlobalTracer
	opentracing.SetGlobalTracer(tracer)

	return func() {
		closer.Close()
	}
}

// InitMetrics configures pkg/metrics
//
// Env Vars:
// * STATSD_ADDR - The host:port of the statsd server.
func InitMetrics() func() {
	if addr := os.Getenv("STATSD_ADDR"); addr != "" {
		metrics.SetEmpireDefaultTags()
		metrics.Reporter, _ = metrics.NewDataDogMetricsReporter(addr)
	}

	return func() {
		metrics.Close()
	}
}

// InitLogger configures a leveled logger.
//
// Env Vars:
// * LOG_LEVEL - The log level
//
// If you want to replace the global default logger:
//	logger.DefaultLogger = InitLogger()
func InitLogger() logger.Logger {
	lvl := logger.ERROR
	if ll := os.Getenv("LOG_LEVEL"); ll != "" {
		lvl = logger.ParseLevel(ll)
	}

	return logger.New(log.New(os.Stdout, "", 0), lvl)
}

// InitReporter configures and returns a reporter.Reporter instance.
//
// Env Vars:
// * ROLLBAR_ACCESS_TOKEN - The Rollbar access token
// * ROLLBAR_ENVIRONMENT  - The Rollbar environment (staging, production)
func InitReporter() reporter.Reporter {
	rbToken := os.Getenv("ROLLBAR_ACCESS_TOKEN")
	rbEnv := os.Getenv("ROLLBAR_ENVIRONMENT")

	rep := reporter.MultiReporter{}

	// Log Reporter, uses package level logger.
	rep = append(rep, reporter.NewLogReporter())

	// Rollbar reporter
	if rbToken != "" && rbEnv != "" {
		rollbar.ConfigureReporter(rbToken, rbEnv)
		rep = append(rep, rollbar.Reporter)
	} else {
		fmt.Println("Rollbar is not configured, skipping Rollbar reporter")
	}

	return rep
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"net/rpc"
	"strconv"
	"time"

	"github.com/coreos/etcd/raft"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/dgraph-io/dgraph/conn"
	"github.com/dgraph-io/dgraph/x"
	"golang.org/x/net/context"
)

var (
	hb    = 1
	pools = make(map[uint64]*conn.Pool)
	glog  = x.Log("RAFT")
	peers = make(map[uint64]string)
)

type node struct {
	id     uint64
	addr   string
	ctx    context.Context
	pstore map[string]string
	store  *raft.MemoryStorage
	cfg    *raft.Config
	raft   raft.Node
	ticker <-chan time.Time
	done   <-chan struct{}
}

type Worker struct {
}

type raftRPC struct {
	Ctx     context.Context
	Message raftpb.Message
}

type helloRPC struct {
	Id   uint64
	Addr string
}

func (w *Worker) Hello(query *conn.Query, reply *conn.Reply) error {
	buf := bytes.NewBuffer(query.Data)
	dec := gob.NewDecoder(buf)
	var v helloRPC
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}

	if _, ok := pools[v.Id]; !ok {
		go connectWith(v.Addr)
	}
	reply.Data = []byte(strconv.Itoa(int(cur_node.id)))

	fmt.Println("In Hello")
	return nil
}

func (w *Worker) JoinCluster(query *conn.Query, reply *conn.Reply) error {
	i, _ := strconv.Atoi(string(query.Data))
	id := uint64(i)
	cur_node.raft.ProposeConfChange(cur_node.ctx, raftpb.ConfChange{
		ID:      id,
		Type:    raftpb.ConfChangeAddNode,
		NodeID:  id,
		Context: []byte(""),
	})
	return nil
}

func serveRequests(irwc io.ReadWriteCloser) {
	for {
		sc := &conn.ServerCodec{
			Rwc: irwc,
		}
		rpc.ServeRequest(sc)
	}
}

func runServer(address string) error {
	ln, err := net.Listen("tcp", address)
	if err != nil {
		glog.Fatalf("While running server: %v", err)
		return err
	}
	glog.WithField("address", ln.Addr()).Info("Worker listening")

	go func() {
		for {
			cxn, err := ln.Accept()
			if err != nil {
				glog.Fatalf("listen(%q): %s\n", address, err)
				return
			}
			glog.WithField("local", cxn.LocalAddr()).
				WithField("remote", cxn.RemoteAddr()).
				Debug("Worker accepted connection")
			go serveRequests(cxn)
		}
	}()
	return nil
}

func newNode(id uint64, addr string, peers []raft.Peer) *node {
	store := raft.NewMemoryStorage()
	n := &node{
		id:    id,
		addr:  addr,
		ctx:   context.TODO(),
		store: store,
		cfg: &raft.Config{
			ID:              id,
			ElectionTick:    5 * hb,
			HeartbeatTick:   hb,
			Storage:         store,
			MaxSizePerMsg:   math.MaxUint16,
			MaxInflightMsgs: 256,
		},
		pstore: make(map[string]string),
		ticker: time.Tick(time.Second),
		done:   make(chan struct{}),
	}

	n.raft = raft.StartNode(n.cfg, peers)
	return n
}

func (n *node) run() {
	for {
		select {
		case <-n.ticker:
			n.raft.Tick()
		case rd := <-n.raft.Ready():
			n.saveToStorage(rd.HardState, rd.Entries, rd.Snapshot)
			n.send(rd.Messages)
			if !raft.IsEmptySnap(rd.Snapshot) {
				n.processSnapshot(rd.Snapshot)
			}
			for _, entry := range rd.CommittedEntries {
				n.process(entry)
				if entry.Type == raftpb.EntryConfChange {
					var cc raftpb.ConfChange
					cc.Unmarshal(entry.Data)
					n.raft.ApplyConfChange(cc)
				}
			}
			n.raft.Advance()
		case <-n.done:
			return
		}
	}
}

func (n *node) saveToStorage(hardState raftpb.HardState, entries []raftpb.Entry, snapshot raftpb.Snapshot) {
	n.store.Append(entries)

	if !raft.IsEmptyHardState(hardState) {
		n.store.SetHardState(hardState)
	}

	if !raft.IsEmptySnap(snapshot) {
		n.store.ApplySnapshot(snapshot)
	}
}

func (n *node) send(messages []raftpb.Message) {
	for _, m := range messages {
		log.Println(raft.DescribeMessage(m, nil))

		// send message to other node
		//nodes[int(m.To)].receive(n.ctx, m)
		sendOverNetwork(n.ctx, m)
	}
}

func sendOverNetwork(ctx context.Context, message raftpb.Message) {
	pool, ok := pools[message.To]
	if !ok {
		glog.WithField("From", cur_node.id).WithField("To", message.To).
			Fatal("Error in making connetions")
	}
	addr := pool.Addr
	fmt.Println(addr)
	query := new(conn.Query)

	var network bytes.Buffer
	gob.Register(ctx)
	enc := gob.NewEncoder(&network)
	err := enc.Encode(raftRPC{ctx, message})
	if err != nil {
		glog.Fatalf("encode:", err)
	}

	query.Data = network.Bytes()
	reply := new(conn.Reply)
	if err := pool.Call("Worker.ReceiveOverNetwork", query, reply); err != nil {
		glog.WithField("call", "Worker.ReceiveOverNetwork").Fatal(err)
	}
	glog.WithField("reply_len", len(reply.Data)).WithField("addr", addr).
		Info("Got reply from server")

}

func (w *Worker) ReceiveOverNetwork(query *conn.Query, reply *conn.Reply) error {
	buf := bytes.NewBuffer(query.Data)
	dec := gob.NewDecoder(buf)
	gob.Register(context.Background())
	var v raftRPC
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}
	cur_node.receive(v.Ctx, v.Message)

	return nil
}

func (n *node) processSnapshot(snapshot raftpb.Snapshot) {
	panic(fmt.Sprintf("Applying snapshot on node %v is not implemented", n.id))
}

func (n *node) process(entry raftpb.Entry) {
	log.Printf("node %v: processing entry: %v\n", n.id, entry)
	if entry.Type == raftpb.EntryNormal && entry.Data != nil {
		parts := bytes.SplitN(entry.Data, []byte(":"), 2)
		n.pstore[string(parts[0])] = string(parts[1])
	}
}

func (n *node) receive(ctx context.Context, message raftpb.Message) {
	n.raft.Step(ctx, message)
}

func connectWith(addr string) uint64 {
	if len(addr) == 0 {
		return 0
	}
	pool := conn.NewPool(addr, 5)
	query := new(conn.Query)
	var network bytes.Buffer
	enc := gob.NewEncoder(&network)
	err := enc.Encode(helloRPC{cur_node.id, *workerPort})
	if err != nil {
		glog.Fatalf("encode:", err)
	}
	query.Data = network.Bytes()

	reply := new(conn.Reply)
	if err := pool.Call("Worker.Hello", query, reply); err != nil {
		glog.WithField("call", "Worker.Hello").Fatal(err)
	}
	i, _ := strconv.Atoi(string(reply.Data))
	glog.WithField("reply", i).WithField("addr", addr).
		Info("Got reply from server")

	pools[uint64(i)] = pool
	peers[uint64(i)] = pool.Addr
	return uint64(i)
}

func proposeJoin(id uint64) {
	pool := pools[id]
	addr := pool.Addr
	fmt.Println(addr)
	query := new(conn.Query)
	query.Data = []byte(strconv.Itoa(int(cur_node.id)))
	reply := new(conn.Reply)
	if err := pool.Call("Worker.JoinCluster", query, reply); err != nil {
		glog.WithField("call", "Worker.JoinCluster").Fatal(err)
	}
	glog.WithField("reply_len", len(reply.Data)).WithField("addr", addr).
		Info("Got reply from server")
}

func (w *Worker) GetPeers(query *conn.Query, reply *conn.Reply) error {
	//gob.Register(pList)
	var network bytes.Buffer
	enc := gob.NewEncoder(&network)
	err := enc.Encode(peers)
	if err != nil {
		glog.Fatalf("encode:", err)
	}

	reply.Data = network.Bytes()
	return nil
}

func getPeerListFrom(id uint64) {
	pool := pools[id]
	addr := pool.Addr
	query := new(conn.Query)
	reply := new(conn.Reply)
	fmt.Println("Got Peer List")
	if err := pool.Call("Worker.GetPeers", query, reply); err != nil {
		glog.WithField("call", "Worker.GetPeers").Fatal(err)
	}
	fmt.Println("Got Peer List")
	glog.WithField("reply_len", len(reply.Data)).WithField("addr", addr).
		Info("Got peerList from server")

	buf := bytes.NewBuffer(reply.Data)
	dec := gob.NewDecoder(buf)
	var v = make(map[uint64]string)
	//gob.Register(pList)
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}
	fmt.Println("Got Peer List")
	updatePeerList(v)
}

func updatePeerList(pl map[uint64]string) {
	for k, v := range pl {
		if _, ok := pools[k]; !ok {
			peers[k] = v
		}
	}
}

func connectWithPeers() {
	for k, v := range peers {
		if _, ok := pools[k]; !ok {
			go connectWith(v)
		}
	}
}

var (
	nodes      = make(map[int]*node)
	w          = new(Worker)
	workerPort = flag.String("workerport", ":12345",
		"Port used by worker for internal communication.")
	instanceIdx = flag.Uint64("idx", 1,
		"raft instance id")
	cluster  = flag.String("clusterIP", "", "IP of a node in cluster")
	cur_node *node
)

func main() {
	flag.Parse()
	cur_node = newNode(*instanceIdx, "", []raft.Peer{{ID: *instanceIdx}})
	if err := rpc.Register(w); err != nil {
		glog.Fatal(err)
	}
	if err := runServer(*workerPort); err != nil {
		glog.Fatal(err)
	}

	peers[*instanceIdx] = *workerPort

	if *cluster != "" {
		i := connectWith(*cluster)
		go cur_node.run()
		getPeerListFrom(i)
		connectWithPeers()
		proposeJoin(i)
	} else {
		go cur_node.run()
	}

	for cur_node.id == 1 && cur_node.raft.Status().Lead != 1 {
		time.Sleep(100 * time.Millisecond)
	}

	fmt.Println("proposal by node ", cur_node.id)
	nodeID := strconv.Itoa(int(cur_node.id))
	cur_node.raft.Propose(cur_node.ctx, []byte("mykey"+nodeID+":myvalue"+nodeID))

	/*
		for cur_node.raft.Status().Lead != 1 {
			time.Sleep(100 * time.Millisecond)
		}
	*/

	count := 0
	for count != 55 {
		count = 0
		fmt.Printf("** Node %v **\n", cur_node.id)
		for k, v := range cur_node.pstore {
			fmt.Printf("%v = %v\n", k, v)
			count += 1
		}
		fmt.Printf("*************\n")
		time.Sleep(1000 * time.Millisecond)
	}

	time.Sleep(1000 * time.Millisecond)
}
<commit_msg>Independent master, fixed blocking network calls<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"net/rpc"
	"strconv"
	"time"

	"github.com/coreos/etcd/raft"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/dgraph-io/dgraph/conn"
	"github.com/dgraph-io/dgraph/x"
	"golang.org/x/net/context"
)

var (
	hb               = 1
	pools            = make(map[uint64]*conn.Pool)
	glog             = x.Log("RAFT")
	peers            = make(map[uint64]string)
	confCount uint64 = 0
)

type node struct {
	id     uint64
	addr   string
	ctx    context.Context
	pstore map[string]string
	store  *raft.MemoryStorage
	cfg    *raft.Config
	raft   raft.Node
	ticker <-chan time.Time
	done   <-chan struct{}
}

type Worker struct {
}

type raftRPC struct {
	Ctx     context.Context
	Message raftpb.Message
}

type helloRPC struct {
	Id   uint64
	Addr string
}

func (w *Worker) Hello(query *conn.Query, reply *conn.Reply) error {
	buf := bytes.NewBuffer(query.Data)
	dec := gob.NewDecoder(buf)
	var v helloRPC
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}

	if _, ok := pools[v.Id]; !ok {
		go connectWith(v.Addr)
	}
	reply.Data = []byte(strconv.Itoa(int(cur_node.id)))

	fmt.Println("In Hello")
	return nil
}

func (w *Worker) JoinCluster(query *conn.Query, reply *conn.Reply) error {
	i, _ := strconv.Atoi(string(query.Data))
	id := uint64(i)
	confCount++
	cur_node.raft.ProposeConfChange(cur_node.ctx, raftpb.ConfChange{
		ID:      confCount,
		Type:    raftpb.ConfChangeAddNode,
		NodeID:  id,
		Context: []byte(""),
	})
	return nil
}

func serveRequests(irwc io.ReadWriteCloser) {
	for {
		sc := &conn.ServerCodec{
			Rwc: irwc,
		}
		rpc.ServeRequest(sc)
	}
}

func runServer(address string) error {
	ln, err := net.Listen("tcp", address)
	if err != nil {
		glog.Fatalf("While running server: %v", err)
		return err
	}
	glog.WithField("address", ln.Addr()).Info("Worker listening")

	go func() {
		for {
			cxn, err := ln.Accept()
			if err != nil {
				glog.Fatalf("listen(%q): %s\n", address, err)
				return
			}
			glog.WithField("local", cxn.LocalAddr()).
				WithField("remote", cxn.RemoteAddr()).
				Debug("Worker accepted connection")
			go serveRequests(cxn)
		}
	}()
	return nil
}

func newNode(id uint64, addr string, peers []raft.Peer) *node {
	store := raft.NewMemoryStorage()
	n := &node{
		id:    id,
		addr:  addr,
		ctx:   context.TODO(),
		store: store,
		cfg: &raft.Config{
			ID:              id,
			ElectionTick:    5 * hb,
			HeartbeatTick:   hb,
			Storage:         store,
			MaxSizePerMsg:   math.MaxUint16,
			MaxInflightMsgs: 256,
		},
		pstore: make(map[string]string),
		ticker: time.Tick(time.Second),
		done:   make(chan struct{}),
	}

	n.raft = raft.StartNode(n.cfg, peers)
	return n
}

func (n *node) run() {
	for {
		select {
		case <-n.ticker:
			n.raft.Tick()
		case rd := <-n.raft.Ready():
			n.saveToStorage(rd.HardState, rd.Entries, rd.Snapshot)
			n.send(rd.Messages)
			if !raft.IsEmptySnap(rd.Snapshot) {
				n.processSnapshot(rd.Snapshot)
			}
			for _, entry := range rd.CommittedEntries {
				n.process(entry)
				if entry.Type == raftpb.EntryConfChange {
					var cc raftpb.ConfChange
					cc.Unmarshal(entry.Data)
					n.raft.ApplyConfChange(cc)
				}
			}
			n.raft.Advance()
		case <-n.done:
			return
		}
	}
}

func (n *node) saveToStorage(hardState raftpb.HardState, entries []raftpb.Entry, snapshot raftpb.Snapshot) {
	n.store.Append(entries)

	if !raft.IsEmptyHardState(hardState) {
		n.store.SetHardState(hardState)
	}

	if !raft.IsEmptySnap(snapshot) {
		n.store.ApplySnapshot(snapshot)
	}
}

func (n *node) send(messages []raftpb.Message) {
	for _, m := range messages {
		log.Println(raft.DescribeMessage(m, nil))

		// send message to other node
		//nodes[int(m.To)].receive(n.ctx, m)
		go sendOverNetwork(n.ctx, m)
	}
}

func sendOverNetwork(ctx context.Context, message raftpb.Message) {
	pool, ok := pools[message.To]
	if !ok {
		glog.WithField("From", cur_node.id).WithField("To", message.To).
			Fatal("Error in making connetions")
	}
	addr := pool.Addr
	fmt.Println(addr)
	query := new(conn.Query)

	var network bytes.Buffer
	gob.Register(ctx)
	enc := gob.NewEncoder(&network)
	err := enc.Encode(raftRPC{ctx, message})
	if err != nil {
		glog.Fatalf("encode:", err)
	}

	query.Data = network.Bytes()
	reply := new(conn.Reply)
	if err := pool.Call("Worker.ReceiveOverNetwork", query, reply); err != nil {
		glog.WithField("call", "Worker.ReceiveOverNetwork").Error(err)
		//cur_node.raft.ReportUnreachable(message.To) // Report to raft cluster that the node is unreachable
		RemoveNodeFromCluster(message.To)
		return
	}
	glog.WithField("reply_len", len(reply.Data)).WithField("addr", addr).
		Info("Got reply from server")

}

func RemoveNodeFromCluster(id uint64) {
	confCount++
	cur_node.raft.ProposeConfChange(cur_node.ctx, raftpb.ConfChange{
		ID:      confCount,
		Type:    raftpb.ConfChangeRemoveNode,
		NodeID:  id,
		Context: []byte(""),
	})
}

func (w *Worker) ReceiveOverNetwork(query *conn.Query, reply *conn.Reply) error {
	buf := bytes.NewBuffer(query.Data)
	dec := gob.NewDecoder(buf)
	gob.Register(context.Background())
	var v raftRPC
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}
	cur_node.receive(v.Ctx, v.Message)

	return nil
}

func (n *node) processSnapshot(snapshot raftpb.Snapshot) {
	panic(fmt.Sprintf("Applying snapshot on node %v is not implemented", n.id))
}

func (n *node) process(entry raftpb.Entry) {
	log.Printf("node %v: processing entry: %v\n", n.id, entry)
	if entry.Type == raftpb.EntryNormal && entry.Data != nil {
		parts := bytes.SplitN(entry.Data, []byte(":"), 2)
		n.pstore[string(parts[0])] = string(parts[1])
	}
}

func (n *node) receive(ctx context.Context, message raftpb.Message) {
	n.raft.Step(ctx, message)
}

func connectWith(addr string) uint64 {
	if len(addr) == 0 {
		return 0
	}
	pool := conn.NewPool(addr, 5)
	query := new(conn.Query)
	var network bytes.Buffer
	enc := gob.NewEncoder(&network)
	err := enc.Encode(helloRPC{cur_node.id, *workerPort})
	if err != nil {
		glog.Fatalf("encode:", err)
	}
	query.Data = network.Bytes()

	reply := new(conn.Reply)
	if err := pool.Call("Worker.Hello", query, reply); err != nil {
		glog.WithField("call", "Worker.Hello").Fatal(err)
	}
	i, _ := strconv.Atoi(string(reply.Data))
	glog.WithField("reply", i).WithField("addr", addr).
		Info("Got reply from server")

	pools[uint64(i)] = pool
	peers[uint64(i)] = pool.Addr
	return uint64(i)
}

func proposeJoin(id uint64) {
	pool := pools[id]
	addr := pool.Addr
	fmt.Println(addr)
	query := new(conn.Query)
	query.Data = []byte(strconv.Itoa(int(cur_node.id)))
	reply := new(conn.Reply)

	co := 0
	for cur_node.raft.Status().Lead != id && co < 3330 {
		glog.Info("Trying to connect with master")
		if err := pool.Call("Worker.JoinCluster", query, reply); err != nil {
			glog.WithField("call", "Worker.JoinCluster").Fatal(err)
		}
		glog.WithField("reply_len", len(reply.Data)).WithField("addr", addr).
			Info("Got reply from server")
		time.Sleep(1000 * time.Millisecond) // sleep for a second and rety joining the cluster
		co++
	}

	if cur_node.raft.Status().Lead != id {
		glog.Fatalf("Unable to joing the cluster")
	}
}
func (w *Worker) GetPeers(query *conn.Query, reply *conn.Reply) error {
	//gob.Register(pList)
	var network bytes.Buffer
	enc := gob.NewEncoder(&network)
	err := enc.Encode(peers)
	if err != nil {
		glog.Fatalf("encode:", err)
	}

	reply.Data = network.Bytes()
	return nil
}

func getPeerListFrom(id uint64) {
	pool := pools[id]
	addr := pool.Addr
	query := new(conn.Query)
	reply := new(conn.Reply)
	fmt.Println("Got Peer List")
	if err := pool.Call("Worker.GetPeers", query, reply); err != nil {
		glog.WithField("call", "Worker.GetPeers").Fatal(err)
	}
	fmt.Println("Got Peer List")
	glog.WithField("reply_len", len(reply.Data)).WithField("addr", addr).
		Info("Got peerList from server")

	buf := bytes.NewBuffer(reply.Data)
	dec := gob.NewDecoder(buf)
	var v = make(map[uint64]string)
	//gob.Register(pList)
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}
	fmt.Println("Got Peer List")
	updatePeerList(v)
}

func updatePeerList(pl map[uint64]string) {
	for k, v := range pl {
		if _, ok := pools[k]; !ok {
			peers[k] = v
		}
	}
}

func connectWithPeers() {
	for k, v := range peers {
		if _, ok := pools[k]; !ok {
			go connectWith(v)
		}
	}
}

func (w *Worker) GetMasterIP(query *conn.Query, reply *conn.Reply) error {
	buf := bytes.NewBuffer(query.Data)
	dec := gob.NewDecoder(buf)
	var v helloRPC
	err := dec.Decode(&v)
	if err != nil {
		glog.Fatal("decode:", err)
	}

	if _, ok := pools[v.Id]; !ok {
		go connectWith(v.Addr)
	}
	reply.Data = []byte(peers[cur_node.raft.Status().Lead])
	fmt.Println("In Hello")
	return nil
}

func getMasterIp(ip string) string {
	if len(ip) == 0 {
		return ""
	}
	pool := conn.NewPool(ip, 5)
	query := new(conn.Query)
	var network bytes.Buffer
	enc := gob.NewEncoder(&network)
	err := enc.Encode(helloRPC{cur_node.id, *workerPort})
	if err != nil {
		glog.Fatalf("encode:", err)
	}
	query.Data = network.Bytes()

	reply := new(conn.Reply)
	if err := pool.Call("Worker.GetMasterIP", query, reply); err != nil {
		glog.WithField("call", "Worker.GetMasterIP").Fatal(err)
	}
	masterIP := string(reply.Data)
	glog.WithField("reply", masterIP).WithField("addr", ip).
		Info("Got reply from server")

	return masterIP
}

var (
	nodes      = make(map[int]*node)
	w          = new(Worker)
	workerPort = flag.String("workerport", ":12345",
		"Port used by worker for internal communication.")
	instanceIdx = flag.Uint64("idx", 1,
		"raft instance id")
	cluster  = flag.String("clusterIP", "", "IP of a node in cluster")
	cur_node *node
)

func main() {
	flag.Parse()
	cur_node = newNode(*instanceIdx, "", []raft.Peer{{ID: *instanceIdx}})
	if err := rpc.Register(w); err != nil {
		glog.Fatal(err)
	}
	if err := runServer(*workerPort); err != nil {
		glog.Fatal(err)
	}

	peers[*instanceIdx] = *workerPort

	if *cluster != "" {
		master_ip := getMasterIp(*cluster)
		master_id := connectWith(master_ip)
		go cur_node.run()
		getPeerListFrom(master_id)
		connectWithPeers()
		proposeJoin(master_id)
	} else {
		go cur_node.run()
		cur_node.raft.Campaign(cur_node.ctx)
	}

	/*
		for cur_node.id == 1 && cur_node.raft.Status().Lead != 1 {
			time.Sleep(100 * time.Millisecond)
		}
	*/

	fmt.Println("proposal by node ", cur_node.id)
	nodeID := strconv.Itoa(int(cur_node.id))
	cur_node.raft.Propose(cur_node.ctx, []byte("mykey"+nodeID+":myvalue"+nodeID))

	/*
		for cur_node.raft.Status().Lead != 1 {
			time.Sleep(100 * time.Millisecond)
		}
	*/

	count := 0
	for count != 55 {
		count = 0
		fmt.Printf("** Node %v **\n", cur_node.id)
		for k, v := range cur_node.pstore {
			fmt.Printf("%v = %v\n", k, v)
			count += 1
		}
		fmt.Printf("*************\n")
		time.Sleep(5 * time.Second)
	}
}
<|endoftext|>