<commit_before>//  Copyright (c) 2015-2017 Rackspace
//
//  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 middleware

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/troubling/hummingbird/client"
	"github.com/troubling/hummingbird/common"
	"github.com/troubling/hummingbird/common/ring"
	"github.com/troubling/hummingbird/common/srv"
	"go.uber.org/zap"
)

var (
	autoCreateAccounts = true // TODO; figure out how to plumb this in as a config
	serverInfo         = make(map[string]interface{})
	sil                sync.Mutex
	excludeHeaders     = []string{
		"X-Account-Sysmeta-",
		"X-Container-Sysmeta-",
		"X-Object-Sysmeta-",
		"X-Object-Transient-Sysmeta-",
		"X-Backend-",
	}
)

func RegisterInfo(name string, data interface{}) {
	sil.Lock()
	defer sil.Unlock()
	serverInfo[name] = data
}

func serverInfoDump() ([]byte, error) {
	sil.Lock()
	defer sil.Unlock()
	data, err := json.Marshal(serverInfo)
	return data, err
}

type AccountInfo struct {
	ContainerCount int64
	ObjectCount    int64
	ObjectBytes    int64
	Metadata       map[string]string
	SysMetadata    map[string]string
}

type AuthorizeFunc func(r *http.Request) bool

type ProxyContextMiddleware struct {
	next              http.Handler
	log               srv.LowLevelLogger
	Cache             ring.MemcacheRing
	proxyDirectClient *client.ProxyDirectClient
}

type ProxyContext struct {
	*ProxyContextMiddleware
	C                client.ProxyClient
	Authorize        AuthorizeFunc
	Logger           srv.LowLevelLogger
	TxId             string
	responseSent     bool
	status           int
	accountInfoCache map[string]*AccountInfo
	depth            int
}

func GetProxyContext(r *http.Request) *ProxyContext {
	if rv := r.Context().Value("proxycontext"); rv != nil {
		return rv.(*ProxyContext)
	}
	return nil
}

func (ctx *ProxyContext) Response() (bool, int) {
	return ctx.responseSent, ctx.status
}

func getPathParts(request *http.Request) (bool, string, string, string) {
	parts := strings.SplitN(request.URL.Path, "/", 5)
	apiRequest := len(parts) > 0 && parts[1] == "v1"
	if len(parts) == 5 {
		return apiRequest, parts[2], parts[3], parts[4]
	} else if len(parts) == 4 {
		return apiRequest, parts[2], parts[3], ""
	} else if len(parts) == 3 {
		return apiRequest, parts[2], "", ""
	}
	return apiRequest, "", "", ""
}

func (ctx *ProxyContext) GetAccountInfo(account string) *AccountInfo {
	var err error
	key := fmt.Sprintf("account/%s", account)
	ai := ctx.accountInfoCache[key]
	if ai == nil {
		if err := ctx.Cache.GetStructured(key, &ai); err != nil {
			ai = nil
		}
	}
	if ai == nil {
		headers, code := ctx.C.HeadAccount(account, nil)
		if code == 404 && autoCreateAccounts {
			ctx.C.PutAccount(account, http.Header{"X-Timestamp": []string{common.GetTimestamp()}})
			headers, code = ctx.C.HeadAccount(account, nil)
		}
		if code/100 != 2 {
			return nil
		}
		ai = &AccountInfo{
			Metadata:    make(map[string]string),
			SysMetadata: make(map[string]string),
		}
		if ai.ContainerCount, err = strconv.ParseInt(headers.Get("X-Account-Container-Count"), 10, 64); err != nil {
			return nil
		}
		if ai.ObjectCount, err = strconv.ParseInt(headers.Get("X-Account-Object-Count"), 10, 64); err != nil {
			return nil
		}
		if ai.ObjectBytes, err = strconv.ParseInt(headers.Get("X-Account-Bytes-Used"), 10, 64); err != nil {
			return nil
		}
		for k := range headers {
			if strings.HasPrefix(k, "X-Account-Meta-") {
				ai.Metadata[k[15:]] = headers.Get(k)
			} else if strings.HasPrefix(k, "X-Account-Sysmeta-") {
				ai.SysMetadata[k[18:]] = headers.Get(k)
			}
		}
		ctx.Cache.Set(key, ai, 30)
	}
	return ai
}

func (ctx *ProxyContext) InvalidateAccountInfo(account string) {
	key := fmt.Sprintf("account/%s", account)
	delete(ctx.accountInfoCache, key)
	ctx.Cache.Delete(key)
}

func (ctx *ProxyContext) Subrequest(writer http.ResponseWriter, req *http.Request, source string) {
	newctx := &ProxyContext{
		ProxyContextMiddleware: ctx.ProxyContextMiddleware,
		Authorize:              ctx.Authorize,
		Logger:                 ctx.Logger.With(zap.String("src", source)),
		C:                      ctx.C,
		TxId:                   ctx.TxId,
		accountInfoCache:       ctx.accountInfoCache,
		responseSent:           false,
		status:                 500,
		depth:                  ctx.depth + 1,
	}
	// TODO: check depth
	newWriter := srv.NewCustomWriter(writer, func(w http.ResponseWriter, status int) int {
		newctx.responseSent = true
		newctx.status = status
		return status
	})
	req.Header.Set("X-Trans-Id", ctx.TxId)
	req.Header.Set("X-Timestamp", common.GetTimestamp())
	newWriter.Header().Set("X-Trans-Id", ctx.TxId)
	req = req.WithContext(context.WithValue(req.Context(), "proxycontext", newctx))
	ctx.next.ServeHTTP(writer, req)
}

func (m *ProxyContextMiddleware) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if !srv.ValidateRequest(writer, request) {
		return
	}

	if request.URL.Path == "/info" {
		if request.Method == "GET" {
			if data, err := serverInfoDump(); err != nil {
				srv.StandardResponse(writer, 500)
			} else {
				writer.WriteHeader(200)
				writer.Header().Set("Content-Type", "application/json")
				writer.Write(data)
			}
			return
		} else if request.Method == "OPTIONS" {
			writer.Header().Set("Allow", "HEAD, GET, OPTIONS")
			writer.WriteHeader(200)
			return
		} else if request.Method == "HEAD" {
			if _, err := serverInfoDump(); err != nil {
				srv.StandardResponse(writer, 500)
			} else {
				writer.WriteHeader(200)
			}
			return
		}
	}

	for k := range request.Header {
		for _, ex := range excludeHeaders {
			if strings.HasPrefix(k, ex) || k == "X-Timestamp" {
				delete(request.Header, k)
			}
		}
	}

	transId := common.GetTransactionId()
	request.Header.Set("X-Trans-Id", transId)
	writer.Header().Set("X-Trans-Id", transId)
	writer.Header().Set("X-Openstack-Request-Id", transId)
	request.Header.Set("X-Timestamp", common.GetTimestamp())
	logr := m.log.With(zap.String("txn", transId))
	ctx := &ProxyContext{
		ProxyContextMiddleware: m,
		Authorize:              nil,
		Logger:                 logr,
		TxId:                   transId,
		responseSent:           false,
		status:                 500,
		accountInfoCache:       make(map[string]*AccountInfo),
		C:                      client.NewProxyClient(m.proxyDirectClient, m.Cache, make(map[string]*client.ContainerInfo)),
	}
	// we'll almost certainly need the AccountInfo and ContainerInfo for the current path, so pre-fetch them in parallel.
	apiRequest, account, container, _ := getPathParts(request)
	if apiRequest && account != "" {
		wg := &sync.WaitGroup{}
		wg.Add(1)
		go func() {
			defer wg.Done()
			if ctx.GetAccountInfo(account) == nil {
				hdr := http.Header{
					"X-Timestamp": []string{common.GetTimestamp()},
				}
				ctx.C.PutAccount(account, hdr) // Auto-create the account if we can't get its info
			}
		}()
		if container != "" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				ctx.C.GetContainerInfo(account, container)
			}()
		}
		wg.Wait()
	}
	newWriter := srv.NewCustomWriter(writer, func(w http.ResponseWriter, status int) int {
		// strip out any bad headers before calling real WriteHeader
		for k := range w.Header() {
			for _, ex := range excludeHeaders {
				if strings.HasPrefix(k, ex) {
					delete(w.Header(), k)
				}
			}
		}
		ctx.responseSent = true
		ctx.status = status
		return status
	})
	request = request.WithContext(context.WithValue(request.Context(), "proxycontext", ctx))
	m.next.ServeHTTP(newWriter, request)
}

func NewContext(mc ring.MemcacheRing, log srv.LowLevelLogger, proxyDirectClient *client.ProxyDirectClient) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return &ProxyContextMiddleware{
			Cache:             mc,
			log:               log,
			next:              next,
			proxyDirectClient: proxyDirectClient,
		}
	}
}
<commit_msg>Addressing review comment<commit_after>//  Copyright (c) 2015-2017 Rackspace
//
//  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 middleware

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/troubling/hummingbird/client"
	"github.com/troubling/hummingbird/common"
	"github.com/troubling/hummingbird/common/ring"
	"github.com/troubling/hummingbird/common/srv"
	"go.uber.org/zap"
)

var (
	autoCreateAccounts = true // TODO; figure out how to plumb this in as a config
	serverInfo         = make(map[string]interface{})
	sil                sync.Mutex
	excludeHeaders     = []string{
		"X-Account-Sysmeta-",
		"X-Container-Sysmeta-",
		"X-Object-Sysmeta-",
		"X-Object-Transient-Sysmeta-",
		"X-Backend-",
	}
)

func RegisterInfo(name string, data interface{}) {
	sil.Lock()
	defer sil.Unlock()
	serverInfo[name] = data
}

func serverInfoDump() ([]byte, error) {
	sil.Lock()
	defer sil.Unlock()
	data, err := json.Marshal(serverInfo)
	return data, err
}

type AccountInfo struct {
	ContainerCount int64
	ObjectCount    int64
	ObjectBytes    int64
	Metadata       map[string]string
	SysMetadata    map[string]string
}

type AuthorizeFunc func(r *http.Request) bool

type ProxyContextMiddleware struct {
	next              http.Handler
	log               srv.LowLevelLogger
	Cache             ring.MemcacheRing
	proxyDirectClient *client.ProxyDirectClient
}

type ProxyContext struct {
	*ProxyContextMiddleware
	C                client.ProxyClient
	Authorize        AuthorizeFunc
	Logger           srv.LowLevelLogger
	TxId             string
	responseSent     bool
	status           int
	accountInfoCache map[string]*AccountInfo
	depth            int
}

func GetProxyContext(r *http.Request) *ProxyContext {
	if rv := r.Context().Value("proxycontext"); rv != nil {
		return rv.(*ProxyContext)
	}
	return nil
}

func (ctx *ProxyContext) Response() (bool, int) {
	return ctx.responseSent, ctx.status
}

func getPathParts(request *http.Request) (bool, string, string, string) {
	parts := strings.SplitN(request.URL.Path, "/", 5)
	apiRequest := len(parts) > 0 && parts[1] == "v1"
	if len(parts) == 5 {
		return apiRequest, parts[2], parts[3], parts[4]
	} else if len(parts) == 4 {
		return apiRequest, parts[2], parts[3], ""
	} else if len(parts) == 3 {
		return apiRequest, parts[2], "", ""
	}
	return apiRequest, "", "", ""
}

func (ctx *ProxyContext) GetAccountInfo(account string) *AccountInfo {
	var err error
	key := fmt.Sprintf("account/%s", account)
	ai := ctx.accountInfoCache[key]
	if ai == nil {
		if err := ctx.Cache.GetStructured(key, &ai); err != nil {
			ai = nil
		}
	}
	if ai == nil {
		headers, code := ctx.C.HeadAccount(account, nil)
		if code == 404 && autoCreateAccounts {
			ctx.C.PutAccount(account, http.Header{"X-Timestamp": []string{common.GetTimestamp()}})
			headers, code = ctx.C.HeadAccount(account, nil)
		}
		if code/100 != 2 {
			return nil
		}
		ai = &AccountInfo{
			Metadata:    make(map[string]string),
			SysMetadata: make(map[string]string),
		}
		if ai.ContainerCount, err = strconv.ParseInt(headers.Get("X-Account-Container-Count"), 10, 64); err != nil {
			return nil
		}
		if ai.ObjectCount, err = strconv.ParseInt(headers.Get("X-Account-Object-Count"), 10, 64); err != nil {
			return nil
		}
		if ai.ObjectBytes, err = strconv.ParseInt(headers.Get("X-Account-Bytes-Used"), 10, 64); err != nil {
			return nil
		}
		for k := range headers {
			if strings.HasPrefix(k, "X-Account-Meta-") {
				ai.Metadata[k[15:]] = headers.Get(k)
			} else if strings.HasPrefix(k, "X-Account-Sysmeta-") {
				ai.SysMetadata[k[18:]] = headers.Get(k)
			}
		}
		ctx.Cache.Set(key, ai, 30)
	}
	return ai
}

func (ctx *ProxyContext) InvalidateAccountInfo(account string) {
	key := fmt.Sprintf("account/%s", account)
	delete(ctx.accountInfoCache, key)
	ctx.Cache.Delete(key)
}

func (ctx *ProxyContext) Subrequest(writer http.ResponseWriter, req *http.Request, source string) {
	newctx := &ProxyContext{
		ProxyContextMiddleware: ctx.ProxyContextMiddleware,
		Authorize:              ctx.Authorize,
		Logger:                 ctx.Logger.With(zap.String("src", source)),
		C:                      ctx.C,
		TxId:                   ctx.TxId,
		accountInfoCache:       ctx.accountInfoCache,
		responseSent:           false,
		status:                 500,
		depth:                  ctx.depth + 1,
	}
	// TODO: check depth
	newWriter := srv.NewCustomWriter(writer, func(w http.ResponseWriter, status int) int {
		newctx.responseSent = true
		newctx.status = status
		return status
	})
	req.Header.Set("X-Trans-Id", ctx.TxId)
	req.Header.Set("X-Timestamp", common.GetTimestamp())
	newWriter.Header().Set("X-Trans-Id", ctx.TxId)
	req = req.WithContext(context.WithValue(req.Context(), "proxycontext", newctx))
	ctx.next.ServeHTTP(writer, req)
}

func (m *ProxyContextMiddleware) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if !srv.ValidateRequest(writer, request) {
		return
	}

	if request.URL.Path == "/info" {
		if request.Method == "GET" {
			if data, err := serverInfoDump(); err != nil {
				srv.StandardResponse(writer, 500)
			} else {
				writer.Header().Set("Content-Type", "application/json")
				writer.WriteHeader(200)
				writer.Write(data)
			}
			return
		} else if request.Method == "OPTIONS" {
			writer.Header().Set("Allow", "HEAD, GET, OPTIONS")
			writer.WriteHeader(200)
			return
		} else if request.Method == "HEAD" {
			if _, err := serverInfoDump(); err != nil {
				srv.StandardResponse(writer, 500)
			} else {
				writer.WriteHeader(200)
			}
			return
		}
	}

	for k := range request.Header {
		for _, ex := range excludeHeaders {
			if strings.HasPrefix(k, ex) || k == "X-Timestamp" {
				delete(request.Header, k)
			}
		}
	}

	transId := common.GetTransactionId()
	request.Header.Set("X-Trans-Id", transId)
	writer.Header().Set("X-Trans-Id", transId)
	writer.Header().Set("X-Openstack-Request-Id", transId)
	request.Header.Set("X-Timestamp", common.GetTimestamp())
	logr := m.log.With(zap.String("txn", transId))
	ctx := &ProxyContext{
		ProxyContextMiddleware: m,
		Authorize:              nil,
		Logger:                 logr,
		TxId:                   transId,
		responseSent:           false,
		status:                 500,
		accountInfoCache:       make(map[string]*AccountInfo),
		C:                      client.NewProxyClient(m.proxyDirectClient, m.Cache, make(map[string]*client.ContainerInfo)),
	}
	// we'll almost certainly need the AccountInfo and ContainerInfo for the current path, so pre-fetch them in parallel.
	apiRequest, account, container, _ := getPathParts(request)
	if apiRequest && account != "" {
		wg := &sync.WaitGroup{}
		wg.Add(1)
		go func() {
			defer wg.Done()
			if ctx.GetAccountInfo(account) == nil {
				hdr := http.Header{
					"X-Timestamp": []string{common.GetTimestamp()},
				}
				ctx.C.PutAccount(account, hdr) // Auto-create the account if we can't get its info
			}
		}()
		if container != "" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				ctx.C.GetContainerInfo(account, container)
			}()
		}
		wg.Wait()
	}
	newWriter := srv.NewCustomWriter(writer, func(w http.ResponseWriter, status int) int {
		// strip out any bad headers before calling real WriteHeader
		for k := range w.Header() {
			for _, ex := range excludeHeaders {
				if strings.HasPrefix(k, ex) {
					delete(w.Header(), k)
				}
			}
		}
		ctx.responseSent = true
		ctx.status = status
		return status
	})
	request = request.WithContext(context.WithValue(request.Context(), "proxycontext", ctx))
	m.next.ServeHTTP(newWriter, request)
}

func NewContext(mc ring.MemcacheRing, log srv.LowLevelLogger, proxyDirectClient *client.ProxyDirectClient) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return &ProxyContextMiddleware{
			Cache:             mc,
			log:               log,
			next:              next,
			proxyDirectClient: proxyDirectClient,
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package workflow_test

import (
	"context"
	"fmt"
	"reflect"
	"strings"
	"testing"
	"time"

	"golang.org/x/build/internal/workflow"
)

func TestTrivial(t *testing.T) {
	wd := workflow.New()
	greeting := wd.Task("echo", echo, wd.Constant("hello world"))
	wd.Output("greeting", greeting)

	w, err := workflow.Start(wd, map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	outputs, err := w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if got, want := outputs["greeting"], "hello world"; got != want {
		t.Errorf("greeting = %q, want %q", got, want)
	}
}

func TestSplitJoin(t *testing.T) {
	wd := workflow.New()
	in := wd.Task("echo", echo, wd.Constant("string #"))
	add1 := wd.Task("add 1", appendInt, in, wd.Constant(1))
	add2 := wd.Task("add 2", appendInt, in, wd.Constant(2))
	both := wd.Slice([]workflow.Value{add1, add2})
	out := wd.Task("join", join, both)
	wd.Output("strings", out)

	w, err := workflow.Start(wd, map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	outputs, err := w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if got, want := outputs["strings"], "string #1,string #2"; got != want {
		t.Errorf("joined output = %q, want %q", got, want)
	}
}

func TestParallelism(t *testing.T) {
	wd := workflow.New()
	out1 := wd.Task("sleep #1", sleep, wd.Constant(100*time.Millisecond))
	out2 := wd.Task("sleep #2", sleep, wd.Constant(100*time.Millisecond))
	wd.Output("out1", out1)
	wd.Output("out2", out2)

	w, err := workflow.Start(wd, map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	start := time.Now()
	_, err = w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if delay := time.Since(start); delay > 150*time.Millisecond {
		t.Errorf("too much time elapsed: %v", delay)
	}
}

func TestParameters(t *testing.T) {
	wd := workflow.New()
	param1 := wd.Parameter("param1")
	param2 := wd.Parameter("param2")
	out1 := wd.Task("echo 1", echo, param1)
	out2 := wd.Task("echo 2", echo, param2)
	wd.Output("out1", out1)
	wd.Output("out2", out2)

	w, err := workflow.Start(wd, map[string]string{"param1": "#1", "param2": "#2"})
	if err != nil {
		t.Fatal(err)
	}
	outputs, err := w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if want := map[string]interface{}{"out1": "#1", "out2": "#2"}; !reflect.DeepEqual(outputs, want) {
		t.Errorf("outputs = %#v, want %#v", outputs, want)
	}
}

func sleep(ctx context.Context, d time.Duration) (struct{}, error) {
	time.Sleep(d)
	return struct{}{}, nil
}

func appendInt(ctx context.Context, s string, i int) (string, error) {
	return fmt.Sprintf("%v%v", s, i), nil
}

func join(ctx context.Context, s []string) (string, error) {
	return strings.Join(s, ","), nil
}

func echo(ctx context.Context, arg string) (string, error) {
	return arg, nil
}

func loggingListener(t *testing.T) func(*workflow.TaskState) {
	return func(st *workflow.TaskState) {
		switch {
		case !st.Finished:
			t.Logf("task %-10v: started", st.Name)
		case st.Error != nil:
			t.Logf("task %-10v: error: %v", st.Name, st.Error)
		default:
			t.Logf("task %-10v: done: %v", st.Name, st.Result)
		}
	}
}
<commit_msg>internal/workflow: fix flaky TestParallelism<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package workflow_test

import (
	"context"
	"fmt"
	"reflect"
	"strings"
	"testing"
	"time"

	"golang.org/x/build/internal/workflow"
)

func TestTrivial(t *testing.T) {
	wd := workflow.New()
	greeting := wd.Task("echo", echo, wd.Constant("hello world"))
	wd.Output("greeting", greeting)

	w, err := workflow.Start(wd, map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	outputs, err := w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if got, want := outputs["greeting"], "hello world"; got != want {
		t.Errorf("greeting = %q, want %q", got, want)
	}
}

func TestSplitJoin(t *testing.T) {
	wd := workflow.New()
	in := wd.Task("echo", echo, wd.Constant("string #"))
	add1 := wd.Task("add 1", appendInt, in, wd.Constant(1))
	add2 := wd.Task("add 2", appendInt, in, wd.Constant(2))
	both := wd.Slice([]workflow.Value{add1, add2})
	out := wd.Task("join", join, both)
	wd.Output("strings", out)

	w, err := workflow.Start(wd, map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	outputs, err := w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if got, want := outputs["strings"], "string #1,string #2"; got != want {
		t.Errorf("joined output = %q, want %q", got, want)
	}
}

func TestParallelism(t *testing.T) {
	// block1 and block2 block until they're both running.
	chan1, chan2 := make(chan bool, 1), make(chan bool, 1)
	block1 := func(ctx context.Context) (string, error) {
		chan1 <- true
		select {
		case <-chan2:
		case <-ctx.Done():
		}
		return "", ctx.Err()
	}
	block2 := func(ctx context.Context) (string, error) {
		chan2 <- true
		select {
		case <-chan1:
		case <-ctx.Done():
		}
		return "", ctx.Err()
	}
	wd := workflow.New()
	out1 := wd.Task("block #1", block1)
	out2 := wd.Task("block #2", block2)
	wd.Output("out1", out1)
	wd.Output("out2", out2)

	w, err := workflow.Start(wd, map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err = w.Run(ctx, loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
}

func TestParameters(t *testing.T) {
	wd := workflow.New()
	param1 := wd.Parameter("param1")
	param2 := wd.Parameter("param2")
	out1 := wd.Task("echo 1", echo, param1)
	out2 := wd.Task("echo 2", echo, param2)
	wd.Output("out1", out1)
	wd.Output("out2", out2)

	w, err := workflow.Start(wd, map[string]string{"param1": "#1", "param2": "#2"})
	if err != nil {
		t.Fatal(err)
	}
	outputs, err := w.Run(context.Background(), loggingListener(t))
	if err != nil {
		t.Fatal(err)
	}
	if want := map[string]interface{}{"out1": "#1", "out2": "#2"}; !reflect.DeepEqual(outputs, want) {
		t.Errorf("outputs = %#v, want %#v", outputs, want)
	}
}

func appendInt(ctx context.Context, s string, i int) (string, error) {
	return fmt.Sprintf("%v%v", s, i), nil
}

func join(ctx context.Context, s []string) (string, error) {
	return strings.Join(s, ","), nil
}

func echo(ctx context.Context, arg string) (string, error) {
	return arg, nil
}

func loggingListener(t *testing.T) func(*workflow.TaskState) {
	return func(st *workflow.TaskState) {
		switch {
		case !st.Finished:
			t.Logf("task %-10v: started", st.Name)
		case st.Error != nil:
			t.Logf("task %-10v: error: %v", st.Name, st.Error)
		default:
			t.Logf("task %-10v: done: %v", st.Name, st.Result)
		}
	}
}
<|endoftext|>


<commit_before>package dataProcess

import (
	"database/sql"
	"github.com/gin-gonic/gin"
	"net/http"
	"time"

	"../autils"
)

type dimRowsInfo struct {
	Domain string `json:"domain"`
	MType  string `json:"type"`
	Num    int    `json:"num"`
}

type dimData struct {
	Columns []tStruct     `json:"columns"`
	Rows    []dimRowsInfo `json:"rows"`
}

// 作弊请求数据处理
func Dimensions(c *gin.Context, db *sql.DB) {
	position := "left"
	cd := dimData{}

	date := c.Query("date")
	if date == "" {
		date = autils.GetCurrentData(time.Now().AddDate(0, 0, -1))
	}

	q, _ := c.Get("conditions")
	sDate := autils.AnaSigleDate(q)
	s := date
	if sDate != "" {
		s = sDate
	}
	cd.Columns = []tStruct{{
		"站点",
		"domain",
		position,
	}, {
		"类型",
		"type",
		position,
	}, {
		"点击量",
		"num",
		position,
	}}

	infos := getDimInfo(db, s)

	cd.Rows = infos

	c.JSON(http.StatusOK, gin.H{
		"status": 0,
		"msg":    "ok",
		"data":   cd,
	})
}

func getDimInfo(db *sql.DB, date string) []dimRowsInfo {
	showText := []string{"", "二跳", "多跳"}
	sqlStr := "select type, url, count from mip_step where date = '" + date + "' order by count desc"
	rows, err := db.Query(sqlStr)
	autils.ErrHadle(err)

	var url string
	var count, dType int
	cri := dimRowsInfo{}
	criArr := []dimRowsInfo{}
	for rows.Next() {
		err := rows.Scan(&dType, &url, &count)
		autils.ErrHadle(err)
		cri.Domain = url
		cri.Num = count
		cri.MType = showText[dType]
		criArr = append(criArr, cri)
	}
	err = rows.Err()
	autils.ErrHadle(err)

	defer rows.Close()
	return criArr
}
<commit_msg>add mutiple page.<commit_after>package dataProcess

import (
	"bytes"
	"database/sql"
	"github.com/gin-gonic/gin"
	"net/http"
	"strconv"
	"time"

	"../autils"
)

type dimRowsInfo struct {
	Domain string `json:"domain"`
	MType  string `json:"type"`
	Num    int    `json:"num"`
}

type dimData struct {
	Columns []tStruct     `json:"columns"`
	Rows    []dimRowsInfo `json:"rows"`
	Total   int           `json:"total"`
}

// 作弊请求数据处理
func Dimensions(c *gin.Context, db *sql.DB) {
	position := "left"
	cd := dimData{}

	start := c.Query("start")
	limit := c.Query("limit")

	date := c.Query("date")
	if date == "" {
		date = autils.GetCurrentData(time.Now().AddDate(0, 0, -1))
	}

	q, _ := c.Get("conditions")
	sDate := autils.AnaSigleDate(q)
	s := date
	if sDate != "" {
		s = sDate
	}
	cd.Columns = []tStruct{{
		"站点",
		"domain",
		position,
	}, {
		"类型",
		"type",
		position,
	}, {
		"点击量",
		"num",
		position,
	}}

	infos, sum := getDimInfo(db, s, start, limit)

	cd.Rows = infos
	cd.Total = sum

	c.JSON(http.StatusOK, gin.H{
		"status": 0,
		"msg":    "ok",
		"data":   cd,
	})
}

func getDimInfo(db *sql.DB, date, start, limit string) ([]dimRowsInfo, int) {
	showText := []string{"", "二跳", "多跳"}

	var sqlStr bytes.Buffer
	sqlStr.WriteString("select type, url, count from mip_step where date = '" + date + "' order by count desc ")

	_, err := strconv.Atoi(limit)
	if err == nil {
		sqlStr.WriteString(" limit " + limit + "")
	}

	_, err = strconv.Atoi(start)
	if err == nil {
		sqlStr.WriteString(" offset " + start + "")
	}

	rows, err := db.Query(sqlStr.String())
	autils.ErrHadle(err)

	var url string
	var count, dType, sum int
	cri := dimRowsInfo{}
	criArr := []dimRowsInfo{}
	for rows.Next() {
		err := rows.Scan(&dType, &url, &count)
		autils.ErrHadle(err)
		cri.Domain = url
		cri.Num = count
		cri.MType = showText[dType]
		criArr = append(criArr, cri)
		sum++
	}
	err = rows.Err()
	autils.ErrHadle(err)

	defer rows.Close()
	return criArr, sum
}
<|endoftext|>


<commit_before>package flect

var pluralRules = []rule{}

// AddPlural adds a rule that will replace the given suffix with the replacement suffix.
func AddPlural(suffix string, repl string) {
	pluralMoot.Lock()
	defer pluralMoot.Unlock()
	pluralRules = append(pluralRules, rule{
		suffix: suffix,
		fn: func(s string) string {
			s = s[:len(s)-len(suffix)]
			return s + repl
		},
	})

	pluralRules = append(pluralRules, rule{
		suffix: repl,
		fn:     noop,
	})
}

var singleToPlural = map[string]string{
	"aircraft":    "aircraft",
	"alias":       "aliases",
	"alumna":      "alumnae",
	"alumnus":     "alumni",
	"analysis":    "analyses",
	"antenna":     "antennas",
	"antithesis":  "antitheses",
	"apex":        "apexes",
	"appendix":    "appendices",
	"axis":        "axes",
	"bacillus":    "bacilli",
	"bacterium":   "bacteria",
	"basis":       "bases",
	"beau":        "beaus",
	"bison":       "bison",
	"bureau":      "bureaus",
	"bus":         "buses",
	"campus":      "campuses",
	"caucus":      "caucuses",
	"child":       "children",
	"château":     "châteaux",
	"circus":      "circuses",
	"codex":       "codices",
	"concerto":    "concertos",
	"corpus":      "corpora",
	"crisis":      "crises",
	"curriculum":  "curriculums",
	"datum":       "data",
	"deer":        "deer",
	"diagnosis":   "diagnoses",
	"die":         "dice",
	"dwarf":       "dwarves",
	"ellipsis":    "ellipses",
	"equipment":   "equipment",
	"erratum":     "errata",
	"faux pas":    "faux pas",
	"fez":         "fezzes",
	"fish":        "fish",
	"focus":       "foci",
	"foo":         "foos",
	"foot":        "feet",
	"formula":     "formulas",
	"fungus":      "fungi",
	"genus":       "genera",
	"goose":       "geese",
	"graffito":    "graffiti",
	"grouse":      "grouse",
	"half":        "halves",
	"halo":        "halos",
	"hoof":        "hooves",
	"human":       "humans",
	"hypothesis":  "hypotheses",
	"index":       "indices",
	"information": "information",
	"jeans":       "jeans",
	"larva":       "larvae",
	"libretto":    "librettos",
	"loaf":        "loaves",
	"locus":       "loci",
	"louse":       "lice",
	"matrix":      "matrices",
	"minutia":     "minutiae",
	"money":       "money",
	"moose":       "moose",
	"mouse":       "mice",
	"nebula":      "nebulae",
	"news":        "news",
	"nucleus":     "nuclei",
	"oasis":       "oases",
	"octopus":     "octopi",
	"offspring":   "offspring",
	"opus":        "opera",
	"ovum":        "ova",
	"ox":          "oxen",
	"parenthesis": "parentheses",
	"phenomenon":  "phenomena",
	"photo":       "photos",
	"phylum":      "phyla",
	"piano":       "pianos",
	"plus":        "pluses",
	"police":      "police",
	"prognosis":   "prognoses",
	"prometheus":  "prometheuses",
	"quiz":        "quizzes",
	"quota":       "quotas",
	"radius":      "radiuses",
	"referendum":  "referendums",
	"ress":        "resses",
	"rice":        "rice",
	"salmon":      "salmon",
	"sex":         "sexes",
	"series":      "series",
	"sheep":       "sheep",
	"shoe":        "shoes",
	"shrimp":      "shrimp",
	"species":     "species",
	"stimulus":    "stimuli",
	"stratum":     "strata",
	"swine":       "swine",
	"syllabus":    "syllabi",
	"symposium":   "symposiums",
	"synapse":     "synapses",
	"synopsis":    "synopses",
	"tableau":     "tableaus",
	"testis":      "testes",
	"thesis":      "theses",
	"thief":       "thieves",
	"tooth":       "teeth",
	"trout":       "trout",
	"tuna":        "tuna",
	"vedalia":     "vedalias",
	"vertebra":    "vertebrae",
	"vertix":      "vertices",
	"vita":        "vitae",
	"vortex":      "vortices",
	"wharf":       "wharves",
	"wife":        "wives",
	"woman":       "women",
	"wolf":        "wolves",
	"you":         "you",
}

var pluralToSingle = map[string]string{}

func init() {
	for k, v := range singleToPlural {
		pluralToSingle[v] = k
	}
}

type singularToPluralSuffix struct {
	singular string
	plural   string
}

var singularToPluralSuffixList = []singularToPluralSuffix{
	{"iterion", "iteria"},
	{"campus", "campuses"},
	{"genera", "genus"},
	{"person", "people"},
	{"phylum", "phyla"},
	{"randum", "randa"},
	{"actus", "acti"},
	{"adium", "adia"},
	{"basis", "basis"},
	{"child", "children"},
	{"chive", "chives"},
	{"focus", "foci"},
	{"hello", "hellos"},
	{"jeans", "jeans"},
	{"louse", "lice"},
	{"media", "media"},
	{"mouse", "mice"},
	{"movie", "movies"},
	{"oasis", "oasis"},
	{"atum", "ata"},
	{"atus", "atuses"},
	{"base", "bases"},
	{"cess", "cesses"},
	{"dium", "diums"},
	{"eses", "esis"},
	{"half", "halves"},
	{"hive", "hives"},
	{"iano", "ianos"},
	{"irus", "iri"},
	{"isis", "ises"},
	{"leus", "li"},
	{"mnus", "mni"},
	{"move", "moves"},
	{"news", "news"},
	{"odex", "odice"},
	{"oose", "eese"},
	{"ouse", "ouses"},
	{"ovum", "ova"},
	{"rion", "ria"},
	{"shoe", "shoes"},
	{"stis", "stes"},
	{"tive", "tives"},
	{"wife", "wives"},
	{"afe", "aves"},
	{"bfe", "bves"},
	{"box", "boxes"},
	{"cfe", "cves"},
	{"dfe", "dves"},
	{"dge", "dges"},
	{"efe", "eves"},
	{"gfe", "gves"},
	{"hfe", "hves"},
	{"ife", "ives"},
	{"itz", "itzes"},
	{"ium", "ia"},
	{"ize", "izes"},
	{"jfe", "jves"},
	{"kfe", "kves"},
	{"man", "men"},
	{"mfe", "mves"},
	{"nfe", "nves"},
	{"nna", "nnas"},
	{"oaf", "oaves"},
	{"oci", "ocus"},
	{"ode", "odes"},
	{"ofe", "oves"},
	{"oot", "eet"},
	{"pfe", "pves"},
	{"pse", "psis"},
	{"qfe", "qves"},
	{"quy", "quies"},
	{"rfe", "rves"},
	{"sfe", "sves"},
	{"tfe", "tves"},
	{"tum", "ta"},
	{"tus", "tuses"},
	{"ufe", "uves"},
	{"ula", "ulae"},
	{"ula", "ulas"},
	{"uli", "ulus"},
	{"use", "uses"},
	{"uss", "usses"},
	{"vfe", "vves"},
	{"wfe", "wves"},
	{"xfe", "xves"},
	{"yfe", "yves"},
	{"you", "you"},
	{"zfe", "zves"},
	{"by", "bies"},
	{"ch", "ches"},
	{"cy", "cies"},
	{"dy", "dies"},
	{"ex", "ices"},
	{"fy", "fies"},
	{"gy", "gies"},
	{"hy", "hies"},
	{"io", "ios"},
	{"jy", "jies"},
	{"ky", "kies"},
	{"lf", "lves"},
	{"ly", "lies"},
	{"my", "mies"},
	{"ny", "nies"},
	{"py", "pies"},
	{"qy", "qies"},
	{"rf", "rves"},
	{"ry", "ries"},
	{"sh", "shes"},
	{"ss", "sses"},
	{"sy", "sies"},
	{"ty", "ties"},
	{"tz", "tzes"},
	{"va", "vae"},
	{"vy", "vies"},
	{"wy", "wies"},
	{"xy", "xies"},
	{"zy", "zies"},
	{"zz", "zzes"},
	{"o", "oes"},
	{"x", "xes"},
}

func init() {
	for _, suffix := range singularToPluralSuffixList {
		AddPlural(suffix.singular, suffix.plural)
		AddSingular(suffix.plural, suffix.singular)
	}
}
<commit_msg>Add "vice", "vices" suffix to fix "device", "service"<commit_after>package flect

var pluralRules = []rule{}

// AddPlural adds a rule that will replace the given suffix with the replacement suffix.
func AddPlural(suffix string, repl string) {
	pluralMoot.Lock()
	defer pluralMoot.Unlock()
	pluralRules = append(pluralRules, rule{
		suffix: suffix,
		fn: func(s string) string {
			s = s[:len(s)-len(suffix)]
			return s + repl
		},
	})

	pluralRules = append(pluralRules, rule{
		suffix: repl,
		fn:     noop,
	})
}

var singleToPlural = map[string]string{
	"aircraft":    "aircraft",
	"alias":       "aliases",
	"alumna":      "alumnae",
	"alumnus":     "alumni",
	"analysis":    "analyses",
	"antenna":     "antennas",
	"antithesis":  "antitheses",
	"apex":        "apexes",
	"appendix":    "appendices",
	"axis":        "axes",
	"bacillus":    "bacilli",
	"bacterium":   "bacteria",
	"basis":       "bases",
	"beau":        "beaus",
	"bison":       "bison",
	"bureau":      "bureaus",
	"bus":         "buses",
	"campus":      "campuses",
	"caucus":      "caucuses",
	"child":       "children",
	"château":     "châteaux",
	"circus":      "circuses",
	"codex":       "codices",
	"concerto":    "concertos",
	"corpus":      "corpora",
	"crisis":      "crises",
	"curriculum":  "curriculums",
	"datum":       "data",
	"deer":        "deer",
	"diagnosis":   "diagnoses",
	"die":         "dice",
	"dwarf":       "dwarves",
	"ellipsis":    "ellipses",
	"equipment":   "equipment",
	"erratum":     "errata",
	"faux pas":    "faux pas",
	"fez":         "fezzes",
	"fish":        "fish",
	"focus":       "foci",
	"foo":         "foos",
	"foot":        "feet",
	"formula":     "formulas",
	"fungus":      "fungi",
	"genus":       "genera",
	"goose":       "geese",
	"graffito":    "graffiti",
	"grouse":      "grouse",
	"half":        "halves",
	"halo":        "halos",
	"hoof":        "hooves",
	"human":       "humans",
	"hypothesis":  "hypotheses",
	"index":       "indices",
	"information": "information",
	"jeans":       "jeans",
	"larva":       "larvae",
	"libretto":    "librettos",
	"loaf":        "loaves",
	"locus":       "loci",
	"louse":       "lice",
	"matrix":      "matrices",
	"minutia":     "minutiae",
	"money":       "money",
	"moose":       "moose",
	"mouse":       "mice",
	"nebula":      "nebulae",
	"news":        "news",
	"nucleus":     "nuclei",
	"oasis":       "oases",
	"octopus":     "octopi",
	"offspring":   "offspring",
	"opus":        "opera",
	"ovum":        "ova",
	"ox":          "oxen",
	"parenthesis": "parentheses",
	"phenomenon":  "phenomena",
	"photo":       "photos",
	"phylum":      "phyla",
	"piano":       "pianos",
	"plus":        "pluses",
	"police":      "police",
	"prognosis":   "prognoses",
	"prometheus":  "prometheuses",
	"quiz":        "quizzes",
	"quota":       "quotas",
	"radius":      "radiuses",
	"referendum":  "referendums",
	"ress":        "resses",
	"rice":        "rice",
	"salmon":      "salmon",
	"sex":         "sexes",
	"series":      "series",
	"sheep":       "sheep",
	"shoe":        "shoes",
	"shrimp":      "shrimp",
	"species":     "species",
	"stimulus":    "stimuli",
	"stratum":     "strata",
	"swine":       "swine",
	"syllabus":    "syllabi",
	"symposium":   "symposiums",
	"synapse":     "synapses",
	"synopsis":    "synopses",
	"tableau":     "tableaus",
	"testis":      "testes",
	"thesis":      "theses",
	"thief":       "thieves",
	"tooth":       "teeth",
	"trout":       "trout",
	"tuna":        "tuna",
	"vedalia":     "vedalias",
	"vertebra":    "vertebrae",
	"vertix":      "vertices",
	"vita":        "vitae",
	"vortex":      "vortices",
	"wharf":       "wharves",
	"wife":        "wives",
	"woman":       "women",
	"wolf":        "wolves",
	"you":         "you",
}

var pluralToSingle = map[string]string{}

func init() {
	for k, v := range singleToPlural {
		pluralToSingle[v] = k
	}
}

type singularToPluralSuffix struct {
	singular string
	plural   string
}

var singularToPluralSuffixList = []singularToPluralSuffix{
	{"iterion", "iteria"},
	{"campus", "campuses"},
	{"genera", "genus"},
	{"person", "people"},
	{"phylum", "phyla"},
	{"randum", "randa"},
	{"actus", "acti"},
	{"adium", "adia"},
	{"basis", "basis"},
	{"child", "children"},
	{"chive", "chives"},
	{"focus", "foci"},
	{"hello", "hellos"},
	{"jeans", "jeans"},
	{"louse", "lice"},
	{"media", "media"},
	{"mouse", "mice"},
	{"movie", "movies"},
	{"oasis", "oasis"},
	{"atum", "ata"},
	{"atus", "atuses"},
	{"base", "bases"},
	{"cess", "cesses"},
	{"dium", "diums"},
	{"eses", "esis"},
	{"half", "halves"},
	{"hive", "hives"},
	{"iano", "ianos"},
	{"irus", "iri"},
	{"isis", "ises"},
	{"leus", "li"},
	{"mnus", "mni"},
	{"move", "moves"},
	{"news", "news"},
	{"odex", "odice"},
	{"oose", "eese"},
	{"ouse", "ouses"},
	{"ovum", "ova"},
	{"rion", "ria"},
	{"shoe", "shoes"},
	{"stis", "stes"},
	{"tive", "tives"},
	{"vice", "vices"},
	{"wife", "wives"},
	{"afe", "aves"},
	{"bfe", "bves"},
	{"box", "boxes"},
	{"cfe", "cves"},
	{"dfe", "dves"},
	{"dge", "dges"},
	{"efe", "eves"},
	{"gfe", "gves"},
	{"hfe", "hves"},
	{"ife", "ives"},
	{"itz", "itzes"},
	{"ium", "ia"},
	{"ize", "izes"},
	{"jfe", "jves"},
	{"kfe", "kves"},
	{"man", "men"},
	{"mfe", "mves"},
	{"nfe", "nves"},
	{"nna", "nnas"},
	{"oaf", "oaves"},
	{"oci", "ocus"},
	{"ode", "odes"},
	{"ofe", "oves"},
	{"oot", "eet"},
	{"pfe", "pves"},
	{"pse", "psis"},
	{"qfe", "qves"},
	{"quy", "quies"},
	{"rfe", "rves"},
	{"sfe", "sves"},
	{"tfe", "tves"},
	{"tum", "ta"},
	{"tus", "tuses"},
	{"ufe", "uves"},
	{"ula", "ulae"},
	{"ula", "ulas"},
	{"uli", "ulus"},
	{"use", "uses"},
	{"uss", "usses"},
	{"vfe", "vves"},
	{"wfe", "wves"},
	{"xfe", "xves"},
	{"yfe", "yves"},
	{"you", "you"},
	{"zfe", "zves"},
	{"by", "bies"},
	{"ch", "ches"},
	{"cy", "cies"},
	{"dy", "dies"},
	{"ex", "ices"},
	{"fy", "fies"},
	{"gy", "gies"},
	{"hy", "hies"},
	{"io", "ios"},
	{"jy", "jies"},
	{"ky", "kies"},
	{"lf", "lves"},
	{"ly", "lies"},
	{"my", "mies"},
	{"ny", "nies"},
	{"py", "pies"},
	{"qy", "qies"},
	{"rf", "rves"},
	{"ry", "ries"},
	{"sh", "shes"},
	{"ss", "sses"},
	{"sy", "sies"},
	{"ty", "ties"},
	{"tz", "tzes"},
	{"va", "vae"},
	{"vy", "vies"},
	{"wy", "wies"},
	{"xy", "xies"},
	{"zy", "zies"},
	{"zz", "zzes"},
	{"o", "oes"},
	{"x", "xes"},
}

func init() {
	for _, suffix := range singularToPluralSuffixList {
		AddPlural(suffix.singular, suffix.plural)
		AddSingular(suffix.plural, suffix.singular)
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gob

import (
	"reflect"
	"testing"
)

type typeT struct {
	id  typeId
	str string
}

var basicTypes = []typeT{
	{tBool, "bool"},
	{tInt, "int"},
	{tUint, "uint"},
	{tFloat, "float"},
	{tBytes, "bytes"},
	{tString, "string"},
}

func getTypeUnlocked(name string, rt reflect.Type) gobType {
	typeLock.Lock()
	defer typeLock.Unlock()
	t, err := getBaseType(name, rt)
	if err != nil {
		panic("getTypeUnlocked: " + err.Error())
	}
	return t
}

// Sanity checks
func TestBasic(t *testing.T) {
	for _, tt := range basicTypes {
		if tt.id.string() != tt.str {
			t.Errorf("checkType: expected %q got %s", tt.str, tt.id.string())
		}
		if tt.id == 0 {
			t.Errorf("id for %q is zero", tt.str)
		}
	}
}

// Reregister some basic types to check registration is idempotent.
func TestReregistration(t *testing.T) {
	newtyp := getTypeUnlocked("int", reflect.TypeOf(int(0)))
	if newtyp != tInt.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("uint", reflect.TypeOf(uint(0)))
	if newtyp != tUint.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("string", reflect.TypeOf("hello"))
	if newtyp != tString.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
}

func TestArrayType(t *testing.T) {
	var a3 [3]int
	a3int := getTypeUnlocked("foo", reflect.TypeOf(a3))
	newa3int := getTypeUnlocked("bar", reflect.TypeOf(a3))
	if a3int != newa3int {
		t.Errorf("second registration of [3]int creates new type")
	}
	var a4 [4]int
	a4int := getTypeUnlocked("goo", reflect.TypeOf(a4))
	if a3int == a4int {
		t.Errorf("registration of [3]int creates same type as [4]int")
	}
	var b3 [3]bool
	a3bool := getTypeUnlocked("", reflect.TypeOf(b3))
	if a3int == a3bool {
		t.Errorf("registration of [3]bool creates same type as [3]int")
	}
	str := a3bool.string()
	expected := "[3]bool"
	if str != expected {
		t.Errorf("array printed as %q; expected %q", str, expected)
	}
}

func TestSliceType(t *testing.T) {
	var s []int
	sint := getTypeUnlocked("slice", reflect.TypeOf(s))
	var news []int
	newsint := getTypeUnlocked("slice1", reflect.TypeOf(news))
	if sint != newsint {
		t.Errorf("second registration of []int creates new type")
	}
	var b []bool
	sbool := getTypeUnlocked("", reflect.TypeOf(b))
	if sbool == sint {
		t.Errorf("registration of []bool creates same type as []int")
	}
	str := sbool.string()
	expected := "[]bool"
	if str != expected {
		t.Errorf("slice printed as %q; expected %q", str, expected)
	}
}

func TestMapType(t *testing.T) {
	var m map[string]int
	mapStringInt := getTypeUnlocked("map", reflect.TypeOf(m))
	var newm map[string]int
	newMapStringInt := getTypeUnlocked("map1", reflect.TypeOf(newm))
	if mapStringInt != newMapStringInt {
		t.Errorf("second registration of map[string]int creates new type")
	}
	var b map[string]bool
	mapStringBool := getTypeUnlocked("", reflect.TypeOf(b))
	if mapStringBool == mapStringInt {
		t.Errorf("registration of map[string]bool creates same type as map[string]int")
	}
	str := mapStringBool.string()
	expected := "map[string]bool"
	if str != expected {
		t.Errorf("map printed as %q; expected %q", str, expected)
	}
}

type Bar struct {
	X string
}

// This structure has pointers and refers to itself, making it a good test case.
type Foo struct {
	A int
	B int32 // will become int
	C string
	D []byte
	E *float64    // will become float64
	F ****float64 // will become float64
	G *Bar
	H *Bar // should not interpolate the definition of Bar again
	I *Foo // will not explode
}

func TestStructType(t *testing.T) {
	sstruct := getTypeUnlocked("Foo", reflect.TypeOf(Foo{}))
	str := sstruct.string()
	// If we can print it correctly, we built it correctly.
	expected := "Foo = struct { A int; B int; C string; D bytes; E float; F float; G Bar = struct { X string; }; H Bar; I Foo; }"
	if str != expected {
		t.Errorf("struct printed as %q; expected %q", str, expected)
	}
}

// Should be OK to register the same type multiple times, as long as they're
// at the same level of indirection.
func TestRegistration(t *testing.T) {
	type T struct{ a int }
	Register(new(T))
	Register(new(T))
}
<commit_msg>encoding/gob: test for type registration name.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gob

import (
	"reflect"
	"testing"
)

type typeT struct {
	id  typeId
	str string
}

var basicTypes = []typeT{
	{tBool, "bool"},
	{tInt, "int"},
	{tUint, "uint"},
	{tFloat, "float"},
	{tBytes, "bytes"},
	{tString, "string"},
}

func getTypeUnlocked(name string, rt reflect.Type) gobType {
	typeLock.Lock()
	defer typeLock.Unlock()
	t, err := getBaseType(name, rt)
	if err != nil {
		panic("getTypeUnlocked: " + err.Error())
	}
	return t
}

// Sanity checks
func TestBasic(t *testing.T) {
	for _, tt := range basicTypes {
		if tt.id.string() != tt.str {
			t.Errorf("checkType: expected %q got %s", tt.str, tt.id.string())
		}
		if tt.id == 0 {
			t.Errorf("id for %q is zero", tt.str)
		}
	}
}

// Reregister some basic types to check registration is idempotent.
func TestReregistration(t *testing.T) {
	newtyp := getTypeUnlocked("int", reflect.TypeOf(int(0)))
	if newtyp != tInt.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("uint", reflect.TypeOf(uint(0)))
	if newtyp != tUint.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("string", reflect.TypeOf("hello"))
	if newtyp != tString.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
}

func TestArrayType(t *testing.T) {
	var a3 [3]int
	a3int := getTypeUnlocked("foo", reflect.TypeOf(a3))
	newa3int := getTypeUnlocked("bar", reflect.TypeOf(a3))
	if a3int != newa3int {
		t.Errorf("second registration of [3]int creates new type")
	}
	var a4 [4]int
	a4int := getTypeUnlocked("goo", reflect.TypeOf(a4))
	if a3int == a4int {
		t.Errorf("registration of [3]int creates same type as [4]int")
	}
	var b3 [3]bool
	a3bool := getTypeUnlocked("", reflect.TypeOf(b3))
	if a3int == a3bool {
		t.Errorf("registration of [3]bool creates same type as [3]int")
	}
	str := a3bool.string()
	expected := "[3]bool"
	if str != expected {
		t.Errorf("array printed as %q; expected %q", str, expected)
	}
}

func TestSliceType(t *testing.T) {
	var s []int
	sint := getTypeUnlocked("slice", reflect.TypeOf(s))
	var news []int
	newsint := getTypeUnlocked("slice1", reflect.TypeOf(news))
	if sint != newsint {
		t.Errorf("second registration of []int creates new type")
	}
	var b []bool
	sbool := getTypeUnlocked("", reflect.TypeOf(b))
	if sbool == sint {
		t.Errorf("registration of []bool creates same type as []int")
	}
	str := sbool.string()
	expected := "[]bool"
	if str != expected {
		t.Errorf("slice printed as %q; expected %q", str, expected)
	}
}

func TestMapType(t *testing.T) {
	var m map[string]int
	mapStringInt := getTypeUnlocked("map", reflect.TypeOf(m))
	var newm map[string]int
	newMapStringInt := getTypeUnlocked("map1", reflect.TypeOf(newm))
	if mapStringInt != newMapStringInt {
		t.Errorf("second registration of map[string]int creates new type")
	}
	var b map[string]bool
	mapStringBool := getTypeUnlocked("", reflect.TypeOf(b))
	if mapStringBool == mapStringInt {
		t.Errorf("registration of map[string]bool creates same type as map[string]int")
	}
	str := mapStringBool.string()
	expected := "map[string]bool"
	if str != expected {
		t.Errorf("map printed as %q; expected %q", str, expected)
	}
}

type Bar struct {
	X string
}

// This structure has pointers and refers to itself, making it a good test case.
type Foo struct {
	A int
	B int32 // will become int
	C string
	D []byte
	E *float64    // will become float64
	F ****float64 // will become float64
	G *Bar
	H *Bar // should not interpolate the definition of Bar again
	I *Foo // will not explode
}

func TestStructType(t *testing.T) {
	sstruct := getTypeUnlocked("Foo", reflect.TypeOf(Foo{}))
	str := sstruct.string()
	// If we can print it correctly, we built it correctly.
	expected := "Foo = struct { A int; B int; C string; D bytes; E float; F float; G Bar = struct { X string; }; H Bar; I Foo; }"
	if str != expected {
		t.Errorf("struct printed as %q; expected %q", str, expected)
	}
}

// Should be OK to register the same type multiple times, as long as they're
// at the same level of indirection.
func TestRegistration(t *testing.T) {
	type T struct{ a int }
	Register(new(T))
	Register(new(T))
}

type N1 struct{}
type N2 struct{}

// See comment in type.go/Register.
func TestRegistrationNaming(t *testing.T) {
	testCases := []struct {
		t    interface{}
		name string
	}{
		{&N1{}, "*gob.N1"},
		{N2{}, "encoding/gob.N2"},
	}

	for _, tc := range testCases {
		Register(tc.t)

		tct := reflect.TypeOf(tc.t)
		if ct := nameToConcreteType[tc.name]; ct != tct {
			t.Errorf("nameToConcreteType[%q] = %v, want %v", tc.name, ct, tct)
		}
		// concreteTypeToName is keyed off the base type.
		if tct.Kind() == reflect.Ptr {
			tct = tct.Elem()
		}
		if n := concreteTypeToName[tct]; n != tc.name {
			t.Errorf("concreteTypeToName[%v] got %v, want %v", tct, n, tc.name)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// 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 files

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"syscall"

	"github.com/coreos/ignition/config/types"
	"github.com/coreos/ignition/internal/exec/stages"
	"github.com/coreos/ignition/internal/exec/util"
	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/resource"
)

const (
	name = "files"
)

var (
	ErrFilesystemUndefined = errors.New("the referenced filesystem was not defined")
)

func init() {
	stages.Register(creator{})
}

type creator struct{}

func (creator) Create(logger *log.Logger, client *resource.HttpClient, root string) stages.Stage {
	return &stage{
		Util: util.Util{
			DestDir: root,
			Logger:  logger,
		},
		client: client,
	}
}

func (creator) Name() string {
	return name
}

type stage struct {
	util.Util

	client *resource.HttpClient
}

func (stage) Name() string {
	return name
}

func (s stage) Run(config types.Config) bool {
	if err := s.createPasswd(config); err != nil {
		s.Logger.Crit("failed to create users/groups: %v", err)
		return false
	}

	if err := s.createFilesystemsFiles(config); err != nil {
		s.Logger.Crit("failed to create files: %v", err)
		return false
	}

	if err := s.createUnits(config); err != nil {
		s.Logger.Crit("failed to create units: %v", err)
		return false
	}

	return true
}

// createFilesystemsFiles creates the files described in config.Storage.Filesystems.
func (s stage) createFilesystemsFiles(config types.Config) error {
	if len(config.Storage.Filesystems) == 0 {
		return nil
	}
	s.Logger.PushPrefix("createFilesystemsFiles")
	defer s.Logger.PopPrefix()

	fileMap, err := s.mapFilesToFilesystems(config)
	if err != nil {
		return err
	}

	for fs, f := range fileMap {
		if err := s.createFiles(fs, f); err != nil {
			return fmt.Errorf("failed to create files: %v", err)
		}
	}

	return nil
}

// mapFilesToFilesystems builds a map of filesystems to files. If multiple
// definitions of the same filesystem are present, only the final definition is
// used.
func (s stage) mapFilesToFilesystems(config types.Config) (map[types.Filesystem][]types.File, error) {
	files := map[string][]types.File{}
	for _, f := range config.Storage.Files {
		files[f.Filesystem] = append(files[f.Filesystem], f)
	}

	filesystems := map[string]types.Filesystem{}
	for _, fs := range config.Storage.Filesystems {
		filesystems[fs.Name] = fs
	}

	fileMap := map[types.Filesystem][]types.File{}
	for fsn, f := range files {
		if fs, ok := filesystems[fsn]; ok {
			fileMap[fs] = append(fileMap[fs], f...)
		} else {
			s.Logger.Crit("the filesystem (%q), was not defined", fsn)
			return nil, ErrFilesystemUndefined
		}
	}

	return fileMap, nil
}

// createFiles creates any files listed for the filesystem in fs.Files.
func (s stage) createFiles(fs types.Filesystem, files []types.File) error {
	s.Logger.PushPrefix("createFiles")
	defer s.Logger.PopPrefix()

	var mnt string
	if fs.Path == nil {
		var err error
		mnt, err = ioutil.TempDir("", "ignition-files")
		if err != nil {
			return fmt.Errorf("failed to create temp directory: %v", err)
		}
		defer os.Remove(mnt)

		dev := string(fs.Mount.Device)
		format := string(fs.Mount.Format)

		if err := s.Logger.LogOp(
			func() error { return syscall.Mount(dev, mnt, format, 0, "") },
			"mounting %q at %q", dev, mnt,
		); err != nil {
			return fmt.Errorf("failed to mount device %q at %q: %v", dev, mnt, err)
		}
		defer s.Logger.LogOp(
			func() error { return syscall.Unmount(mnt, 0) },
			"unmounting %q at %q", dev, mnt,
		)
	} else {
		mnt = string(*fs.Path)
	}

	u := util.Util{
		Logger:  s.Logger,
		DestDir: mnt,
	}
	for _, f := range files {
		file := util.RenderFile(s.Logger, s.client, f)
		if file == nil {
			return fmt.Errorf("failed to resolve file %q", f.Path)
		}

		if err := s.Logger.LogOp(
			func() error { return u.WriteFile(file) },
			"writing file %q", string(f.Path),
		); err != nil {
			return fmt.Errorf("failed to create file %q: %v", file.Path, err)
		}
	}

	return nil
}

// createUnits creates the units listed under systemd.units and networkd.units.
func (s stage) createUnits(config types.Config) error {
	for _, unit := range config.Systemd.Units {
		if err := s.writeSystemdUnit(unit); err != nil {
			return err
		}
		if unit.Enable {
			if err := s.Logger.LogOp(
				func() error { return s.EnableUnit(unit) },
				"enabling unit %q", unit.Name,
			); err != nil {
				return err
			}
		}
		if unit.Mask {
			if err := s.Logger.LogOp(
				func() error { return s.MaskUnit(unit) },
				"masking unit %q", unit.Name,
			); err != nil {
				return err
			}
		}
	}
	for _, unit := range config.Networkd.Units {
		if err := s.writeNetworkdUnit(unit); err != nil {
			return err
		}
	}
	return nil
}

// writeSystemdUnit creates the specified unit and any dropins for that unit.
// If the contents of the unit or are empty, the unit is not created. The same
// applies to the unit's dropins.
func (s stage) writeSystemdUnit(unit types.SystemdUnit) error {
	return s.Logger.LogOp(func() error {
		for _, dropin := range unit.DropIns {
			if dropin.Contents == "" {
				continue
			}

			f := util.FileFromUnitDropin(unit, dropin)
			if err := s.Logger.LogOp(
				func() error { return s.WriteFile(f) },
				"writing dropin %q at %q", dropin.Name, f.Path,
			); err != nil {
				return err
			}
		}

		if unit.Contents == "" {
			return nil
		}

		f := util.FileFromSystemdUnit(unit)
		if err := s.Logger.LogOp(
			func() error { return s.WriteFile(f) },
			"writing unit %q at %q", unit.Name, f.Path,
		); err != nil {
			return err
		}

		return nil
	}, "writing unit %q", unit.Name)
}

// writeNetworkdUnit creates the specified unit. If the contents of the unit or
// are empty, the unit is not created.
func (s stage) writeNetworkdUnit(unit types.NetworkdUnit) error {
	return s.Logger.LogOp(func() error {
		if unit.Contents == "" {
			return nil
		}

		f := util.FileFromNetworkdUnit(unit)
		if err := s.Logger.LogOp(
			func() error { return s.WriteFile(f) },
			"writing unit %q at %q", unit.Name, f.Path,
		); err != nil {
			return err
		}

		return nil
	}, "writing unit %q", unit.Name)
}

// createPasswd creates the users and groups as described in config.Passwd.
func (s stage) createPasswd(config types.Config) error {
	if err := s.createGroups(config); err != nil {
		return fmt.Errorf("failed to create groups: %v", err)
	}

	if err := s.createUsers(config); err != nil {
		return fmt.Errorf("failed to create users: %v", err)
	}

	return nil
}

// createUsers creates the users as described in config.Passwd.Users.
func (s stage) createUsers(config types.Config) error {
	if len(config.Passwd.Users) == 0 {
		return nil
	}
	s.Logger.PushPrefix("createUsers")
	defer s.Logger.PopPrefix()

	for _, u := range config.Passwd.Users {
		if err := s.CreateUser(u); err != nil {
			return fmt.Errorf("failed to create user %q: %v",
				u.Name, err)
		}

		if err := s.SetPasswordHash(u); err != nil {
			return fmt.Errorf("failed to set password for %q: %v",
				u.Name, err)
		}

		if err := s.AuthorizeSSHKeys(u); err != nil {
			return fmt.Errorf("failed to add keys to user %q: %v",
				u.Name, err)
		}
	}

	return nil
}

// createGroups creates the users as described in config.Passwd.Groups.
func (s stage) createGroups(config types.Config) error {
	if len(config.Passwd.Groups) == 0 {
		return nil
	}
	s.Logger.PushPrefix("createGroups")
	defer s.Logger.PopPrefix()

	for _, g := range config.Passwd.Groups {
		if err := s.CreateGroup(g); err != nil {
			return fmt.Errorf("failed to create group %q: %v",
				g.Name, err)
		}
	}

	return nil
}
<commit_msg>exec/stages/files: improve logging around units<commit_after>// Copyright 2015 CoreOS, Inc.
//
// 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 files

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"syscall"

	"github.com/coreos/ignition/config/types"
	"github.com/coreos/ignition/internal/exec/stages"
	"github.com/coreos/ignition/internal/exec/util"
	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/resource"
)

const (
	name = "files"
)

var (
	ErrFilesystemUndefined = errors.New("the referenced filesystem was not defined")
)

func init() {
	stages.Register(creator{})
}

type creator struct{}

func (creator) Create(logger *log.Logger, client *resource.HttpClient, root string) stages.Stage {
	return &stage{
		Util: util.Util{
			DestDir: root,
			Logger:  logger,
		},
		client: client,
	}
}

func (creator) Name() string {
	return name
}

type stage struct {
	util.Util

	client *resource.HttpClient
}

func (stage) Name() string {
	return name
}

func (s stage) Run(config types.Config) bool {
	if err := s.createPasswd(config); err != nil {
		s.Logger.Crit("failed to create users/groups: %v", err)
		return false
	}

	if err := s.createFilesystemsFiles(config); err != nil {
		s.Logger.Crit("failed to create files: %v", err)
		return false
	}

	if err := s.createUnits(config); err != nil {
		s.Logger.Crit("failed to create units: %v", err)
		return false
	}

	return true
}

// createFilesystemsFiles creates the files described in config.Storage.Filesystems.
func (s stage) createFilesystemsFiles(config types.Config) error {
	if len(config.Storage.Filesystems) == 0 {
		return nil
	}
	s.Logger.PushPrefix("createFilesystemsFiles")
	defer s.Logger.PopPrefix()

	fileMap, err := s.mapFilesToFilesystems(config)
	if err != nil {
		return err
	}

	for fs, f := range fileMap {
		if err := s.createFiles(fs, f); err != nil {
			return fmt.Errorf("failed to create files: %v", err)
		}
	}

	return nil
}

// mapFilesToFilesystems builds a map of filesystems to files. If multiple
// definitions of the same filesystem are present, only the final definition is
// used.
func (s stage) mapFilesToFilesystems(config types.Config) (map[types.Filesystem][]types.File, error) {
	files := map[string][]types.File{}
	for _, f := range config.Storage.Files {
		files[f.Filesystem] = append(files[f.Filesystem], f)
	}

	filesystems := map[string]types.Filesystem{}
	for _, fs := range config.Storage.Filesystems {
		filesystems[fs.Name] = fs
	}

	fileMap := map[types.Filesystem][]types.File{}
	for fsn, f := range files {
		if fs, ok := filesystems[fsn]; ok {
			fileMap[fs] = append(fileMap[fs], f...)
		} else {
			s.Logger.Crit("the filesystem (%q), was not defined", fsn)
			return nil, ErrFilesystemUndefined
		}
	}

	return fileMap, nil
}

// createFiles creates any files listed for the filesystem in fs.Files.
func (s stage) createFiles(fs types.Filesystem, files []types.File) error {
	s.Logger.PushPrefix("createFiles")
	defer s.Logger.PopPrefix()

	var mnt string
	if fs.Path == nil {
		var err error
		mnt, err = ioutil.TempDir("", "ignition-files")
		if err != nil {
			return fmt.Errorf("failed to create temp directory: %v", err)
		}
		defer os.Remove(mnt)

		dev := string(fs.Mount.Device)
		format := string(fs.Mount.Format)

		if err := s.Logger.LogOp(
			func() error { return syscall.Mount(dev, mnt, format, 0, "") },
			"mounting %q at %q", dev, mnt,
		); err != nil {
			return fmt.Errorf("failed to mount device %q at %q: %v", dev, mnt, err)
		}
		defer s.Logger.LogOp(
			func() error { return syscall.Unmount(mnt, 0) },
			"unmounting %q at %q", dev, mnt,
		)
	} else {
		mnt = string(*fs.Path)
	}

	u := util.Util{
		Logger:  s.Logger,
		DestDir: mnt,
	}
	for _, f := range files {
		file := util.RenderFile(s.Logger, s.client, f)
		if file == nil {
			return fmt.Errorf("failed to resolve file %q", f.Path)
		}

		if err := s.Logger.LogOp(
			func() error { return u.WriteFile(file) },
			"writing file %q", string(f.Path),
		); err != nil {
			return fmt.Errorf("failed to create file %q: %v", file.Path, err)
		}
	}

	return nil
}

// createUnits creates the units listed under systemd.units and networkd.units.
func (s stage) createUnits(config types.Config) error {
	for _, unit := range config.Systemd.Units {
		if err := s.writeSystemdUnit(unit); err != nil {
			return err
		}
		if unit.Enable {
			if err := s.Logger.LogOp(
				func() error { return s.EnableUnit(unit) },
				"enabling unit %q", unit.Name,
			); err != nil {
				return err
			}
		}
		if unit.Mask {
			if err := s.Logger.LogOp(
				func() error { return s.MaskUnit(unit) },
				"masking unit %q", unit.Name,
			); err != nil {
				return err
			}
		}
	}
	for _, unit := range config.Networkd.Units {
		if err := s.writeNetworkdUnit(unit); err != nil {
			return err
		}
	}
	return nil
}

// writeSystemdUnit creates the specified unit and any dropins for that unit.
// If the contents of the unit or are empty, the unit is not created. The same
// applies to the unit's dropins.
func (s stage) writeSystemdUnit(unit types.SystemdUnit) error {
	return s.Logger.LogOp(func() error {
		for _, dropin := range unit.DropIns {
			if dropin.Contents == "" {
				continue
			}

			f := util.FileFromUnitDropin(unit, dropin)
			if err := s.Logger.LogOp(
				func() error { return s.WriteFile(f) },
				"writing drop-in %q at %q", dropin.Name, f.Path,
			); err != nil {
				return err
			}
		}

		if unit.Contents == "" {
			return nil
		}

		f := util.FileFromSystemdUnit(unit)
		if err := s.Logger.LogOp(
			func() error { return s.WriteFile(f) },
			"writing unit %q at %q", unit.Name, f.Path,
		); err != nil {
			return err
		}

		return nil
	}, "processing unit %q", unit.Name)
}

// writeNetworkdUnit creates the specified unit. If the contents of the unit or
// are empty, the unit is not created.
func (s stage) writeNetworkdUnit(unit types.NetworkdUnit) error {
	return s.Logger.LogOp(func() error {
		if unit.Contents == "" {
			return nil
		}

		f := util.FileFromNetworkdUnit(unit)
		if err := s.Logger.LogOp(
			func() error { return s.WriteFile(f) },
			"writing unit %q at %q", unit.Name, f.Path,
		); err != nil {
			return err
		}

		return nil
	}, "processing unit %q", unit.Name)
}

// createPasswd creates the users and groups as described in config.Passwd.
func (s stage) createPasswd(config types.Config) error {
	if err := s.createGroups(config); err != nil {
		return fmt.Errorf("failed to create groups: %v", err)
	}

	if err := s.createUsers(config); err != nil {
		return fmt.Errorf("failed to create users: %v", err)
	}

	return nil
}

// createUsers creates the users as described in config.Passwd.Users.
func (s stage) createUsers(config types.Config) error {
	if len(config.Passwd.Users) == 0 {
		return nil
	}
	s.Logger.PushPrefix("createUsers")
	defer s.Logger.PopPrefix()

	for _, u := range config.Passwd.Users {
		if err := s.CreateUser(u); err != nil {
			return fmt.Errorf("failed to create user %q: %v",
				u.Name, err)
		}

		if err := s.SetPasswordHash(u); err != nil {
			return fmt.Errorf("failed to set password for %q: %v",
				u.Name, err)
		}

		if err := s.AuthorizeSSHKeys(u); err != nil {
			return fmt.Errorf("failed to add keys to user %q: %v",
				u.Name, err)
		}
	}

	return nil
}

// createGroups creates the users as described in config.Passwd.Groups.
func (s stage) createGroups(config types.Config) error {
	if len(config.Passwd.Groups) == 0 {
		return nil
	}
	s.Logger.PushPrefix("createGroups")
	defer s.Logger.PopPrefix()

	for _, g := range config.Passwd.Groups {
		if err := s.CreateGroup(g); err != nil {
			return fmt.Errorf("failed to create group %q: %v",
				g.Name, err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package repository

import (
	"context"
	"sync"

	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/ui/progress"
)

// MasterIndex is a collection of indexes and IDs of chunks that are in the process of being saved.
type MasterIndex struct {
	idx          []*Index
	pendingBlobs restic.BlobSet
	idxMutex     sync.RWMutex
}

// NewMasterIndex creates a new master index.
func NewMasterIndex() *MasterIndex {
	// Always add an empty final index, such that MergeFinalIndexes can merge into this.
	// Note that removing this index could lead to a race condition in the rare
	// sitation that only two indexes exist which are saved and merged concurrently.
	idx := []*Index{NewIndex()}
	idx[0].Finalize()
	return &MasterIndex{idx: idx, pendingBlobs: restic.NewBlobSet()}
}

// Lookup queries all known Indexes for the ID and returns all matches.
func (mi *MasterIndex) Lookup(id restic.ID, tpe restic.BlobType) (blobs []restic.PackedBlob) {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	for _, idx := range mi.idx {
		blobs = idx.Lookup(id, tpe, blobs)
	}

	return blobs
}

// LookupSize queries all known Indexes for the ID and returns the first match.
func (mi *MasterIndex) LookupSize(id restic.ID, tpe restic.BlobType) (uint, bool) {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	for _, idx := range mi.idx {
		if size, found := idx.LookupSize(id, tpe); found {
			return size, found
		}
	}

	return 0, false
}

// AddPending adds a given blob to list of pending Blobs
// Before doing so it checks if this blob is already known.
// Returns true if adding was successful and false if the blob
// was already known
func (mi *MasterIndex) addPending(id restic.ID, tpe restic.BlobType) bool {

	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	// Check if blob is pending or in index
	if mi.pendingBlobs.Has(restic.BlobHandle{ID: id, Type: tpe}) {
		return false
	}

	for _, idx := range mi.idx {
		if idx.Has(id, tpe) {
			return false
		}
	}

	// really not known -> insert
	mi.pendingBlobs.Insert(restic.BlobHandle{ID: id, Type: tpe})
	return true
}

// Has queries all known Indexes for the ID and returns the first match.
// Also returns true if the ID is pending.
func (mi *MasterIndex) Has(id restic.ID, tpe restic.BlobType) bool {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	// also return true if blob is pending
	if mi.pendingBlobs.Has(restic.BlobHandle{ID: id, Type: tpe}) {
		return true
	}

	for _, idx := range mi.idx {
		if idx.Has(id, tpe) {
			return true
		}
	}

	return false
}

// Packs returns all packs that are covered by the index.
func (mi *MasterIndex) Packs() restic.IDSet {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	packs := restic.NewIDSet()
	for _, idx := range mi.idx {
		packs.Merge(idx.Packs())
	}

	return packs
}

// Count returns the number of blobs of type t in the index.
func (mi *MasterIndex) Count(t restic.BlobType) (n uint) {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	var sum uint
	for _, idx := range mi.idx {
		sum += idx.Count(t)
	}

	return sum
}

// Insert adds a new index to the MasterIndex.
func (mi *MasterIndex) Insert(idx *Index) {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	mi.idx = append(mi.idx, idx)
}

// StorePack remembers the id and pack in the index.
func (mi *MasterIndex) StorePack(id restic.ID, blobs []restic.Blob) {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	// delete blobs from pending
	for _, blob := range blobs {
		mi.pendingBlobs.Delete(restic.BlobHandle{Type: blob.Type, ID: blob.ID})
	}

	for _, idx := range mi.idx {
		if !idx.Final() {
			idx.StorePack(id, blobs)
			return
		}
	}

	newIdx := NewIndex()
	newIdx.StorePack(id, blobs)
	mi.idx = append(mi.idx, newIdx)
}

// FinalizeNotFinalIndexes finalizes all indexes that
// have not yet been saved and returns that list
func (mi *MasterIndex) FinalizeNotFinalIndexes() []*Index {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	var list []*Index

	for _, idx := range mi.idx {
		if !idx.Final() {
			idx.Finalize()
			list = append(list, idx)
		}
	}

	debug.Log("return %d indexes", len(list))
	return list
}

// FinalizeFullIndexes finalizes all indexes that are full and returns that list.
func (mi *MasterIndex) FinalizeFullIndexes() []*Index {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	var list []*Index

	debug.Log("checking %d indexes", len(mi.idx))
	for _, idx := range mi.idx {
		if idx.Final() {
			debug.Log("index %p is final", idx)
			continue
		}

		if IndexFull(idx) {
			debug.Log("index %p is full", idx)
			idx.Finalize()
			list = append(list, idx)
		} else {
			debug.Log("index %p not full", idx)
		}
	}

	debug.Log("return %d indexes", len(list))
	return list
}

// All returns all indexes.
func (mi *MasterIndex) All() []*Index {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	return mi.idx
}

// Each returns a channel that yields all blobs known to the index. When the
// context is cancelled, the background goroutine terminates. This blocks any
// modification of the index.
func (mi *MasterIndex) Each(ctx context.Context) <-chan restic.PackedBlob {
	mi.idxMutex.RLock()

	ch := make(chan restic.PackedBlob)

	go func() {
		defer mi.idxMutex.RUnlock()
		defer func() {
			close(ch)
		}()

		for _, idx := range mi.idx {
			idxCh := idx.Each(ctx)
			for pb := range idxCh {
				select {
				case <-ctx.Done():
					return
				case ch <- pb:
				}
			}
		}
	}()

	return ch
}

// MergeFinalIndexes merges all final indexes together.
// After calling, there will be only one big final index in MasterIndex
// containing all final index contents.
// Indexes that are not final are left untouched.
// This merging can only be called after all index files are loaded - as
// removing of superseded index contents is only possible for unmerged indexes.
func (mi *MasterIndex) MergeFinalIndexes() {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	// The first index is always final and the one to merge into
	newIdx := mi.idx[:1]
	for i := 1; i < len(mi.idx); i++ {
		idx := mi.idx[i]
		// clear reference in masterindex as it may become stale
		mi.idx[i] = nil
		if !idx.Final() {
			newIdx = append(newIdx, idx)
		} else {
			mi.idx[0].merge(idx)
		}
	}
	mi.idx = newIdx
}

// Save saves all known indexes to index files, leaving out any
// packs whose ID is contained in packBlacklist. The new index contains the IDs
// of all known indexes in the "supersedes" field. The IDs are also returned in
// the IDSet obsolete
// After calling this function, you should remove the obsolete index files.
func (mi *MasterIndex) Save(ctx context.Context, repo restic.Repository, packBlacklist restic.IDSet, extraObsolete restic.IDs, p *progress.Counter) (obsolete restic.IDSet, err error) {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	debug.Log("start rebuilding index of %d indexes, pack blacklist: %v", len(mi.idx), packBlacklist)

	newIndex := NewIndex()
	obsolete = restic.NewIDSet()

	for i, idx := range mi.idx {
		if idx.Final() {
			ids, err := idx.IDs()
			if err != nil {
				debug.Log("index %d does not have an ID: %v", err)
				return nil, err
			}

			debug.Log("adding index ids %v to supersedes field", ids)

			err = newIndex.AddToSupersedes(ids...)
			if err != nil {
				return nil, err
			}
			obsolete.Merge(restic.NewIDSet(ids...))
		} else {
			debug.Log("index %d isn't final, don't add to supersedes field", i)
		}

		debug.Log("adding index %d", i)

		for pbs := range idx.EachByPack(ctx, packBlacklist) {
			newIndex.StorePack(pbs.packID, pbs.blobs)
			p.Add(1)
			if IndexFull(newIndex) {
				newIndex.Finalize()
				if _, err := SaveIndex(ctx, repo, newIndex); err != nil {
					return nil, err
				}
				newIndex = NewIndex()
			}
		}
	}

	err = newIndex.AddToSupersedes(extraObsolete...)
	if err != nil {
		return nil, err
	}
	obsolete.Merge(restic.NewIDSet(extraObsolete...))

	newIndex.Finalize()
	if _, err := SaveIndex(ctx, repo, newIndex); err != nil {
		return nil, err
	}

	return
}
<commit_msg>Parallelize MasterIndex.Save()<commit_after>package repository

import (
	"context"
	"sync"

	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/ui/progress"
	"golang.org/x/sync/errgroup"
)

// MasterIndex is a collection of indexes and IDs of chunks that are in the process of being saved.
type MasterIndex struct {
	idx          []*Index
	pendingBlobs restic.BlobSet
	idxMutex     sync.RWMutex
}

// NewMasterIndex creates a new master index.
func NewMasterIndex() *MasterIndex {
	// Always add an empty final index, such that MergeFinalIndexes can merge into this.
	// Note that removing this index could lead to a race condition in the rare
	// sitation that only two indexes exist which are saved and merged concurrently.
	idx := []*Index{NewIndex()}
	idx[0].Finalize()
	return &MasterIndex{idx: idx, pendingBlobs: restic.NewBlobSet()}
}

// Lookup queries all known Indexes for the ID and returns all matches.
func (mi *MasterIndex) Lookup(id restic.ID, tpe restic.BlobType) (blobs []restic.PackedBlob) {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	for _, idx := range mi.idx {
		blobs = idx.Lookup(id, tpe, blobs)
	}

	return blobs
}

// LookupSize queries all known Indexes for the ID and returns the first match.
func (mi *MasterIndex) LookupSize(id restic.ID, tpe restic.BlobType) (uint, bool) {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	for _, idx := range mi.idx {
		if size, found := idx.LookupSize(id, tpe); found {
			return size, found
		}
	}

	return 0, false
}

// AddPending adds a given blob to list of pending Blobs
// Before doing so it checks if this blob is already known.
// Returns true if adding was successful and false if the blob
// was already known
func (mi *MasterIndex) addPending(id restic.ID, tpe restic.BlobType) bool {

	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	// Check if blob is pending or in index
	if mi.pendingBlobs.Has(restic.BlobHandle{ID: id, Type: tpe}) {
		return false
	}

	for _, idx := range mi.idx {
		if idx.Has(id, tpe) {
			return false
		}
	}

	// really not known -> insert
	mi.pendingBlobs.Insert(restic.BlobHandle{ID: id, Type: tpe})
	return true
}

// Has queries all known Indexes for the ID and returns the first match.
// Also returns true if the ID is pending.
func (mi *MasterIndex) Has(id restic.ID, tpe restic.BlobType) bool {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	// also return true if blob is pending
	if mi.pendingBlobs.Has(restic.BlobHandle{ID: id, Type: tpe}) {
		return true
	}

	for _, idx := range mi.idx {
		if idx.Has(id, tpe) {
			return true
		}
	}

	return false
}

// Packs returns all packs that are covered by the index.
func (mi *MasterIndex) Packs() restic.IDSet {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	packs := restic.NewIDSet()
	for _, idx := range mi.idx {
		packs.Merge(idx.Packs())
	}

	return packs
}

// Count returns the number of blobs of type t in the index.
func (mi *MasterIndex) Count(t restic.BlobType) (n uint) {
	mi.idxMutex.RLock()
	defer mi.idxMutex.RUnlock()

	var sum uint
	for _, idx := range mi.idx {
		sum += idx.Count(t)
	}

	return sum
}

// Insert adds a new index to the MasterIndex.
func (mi *MasterIndex) Insert(idx *Index) {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	mi.idx = append(mi.idx, idx)
}

// StorePack remembers the id and pack in the index.
func (mi *MasterIndex) StorePack(id restic.ID, blobs []restic.Blob) {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	// delete blobs from pending
	for _, blob := range blobs {
		mi.pendingBlobs.Delete(restic.BlobHandle{Type: blob.Type, ID: blob.ID})
	}

	for _, idx := range mi.idx {
		if !idx.Final() {
			idx.StorePack(id, blobs)
			return
		}
	}

	newIdx := NewIndex()
	newIdx.StorePack(id, blobs)
	mi.idx = append(mi.idx, newIdx)
}

// FinalizeNotFinalIndexes finalizes all indexes that
// have not yet been saved and returns that list
func (mi *MasterIndex) FinalizeNotFinalIndexes() []*Index {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	var list []*Index

	for _, idx := range mi.idx {
		if !idx.Final() {
			idx.Finalize()
			list = append(list, idx)
		}
	}

	debug.Log("return %d indexes", len(list))
	return list
}

// FinalizeFullIndexes finalizes all indexes that are full and returns that list.
func (mi *MasterIndex) FinalizeFullIndexes() []*Index {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	var list []*Index

	debug.Log("checking %d indexes", len(mi.idx))
	for _, idx := range mi.idx {
		if idx.Final() {
			debug.Log("index %p is final", idx)
			continue
		}

		if IndexFull(idx) {
			debug.Log("index %p is full", idx)
			idx.Finalize()
			list = append(list, idx)
		} else {
			debug.Log("index %p not full", idx)
		}
	}

	debug.Log("return %d indexes", len(list))
	return list
}

// All returns all indexes.
func (mi *MasterIndex) All() []*Index {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	return mi.idx
}

// Each returns a channel that yields all blobs known to the index. When the
// context is cancelled, the background goroutine terminates. This blocks any
// modification of the index.
func (mi *MasterIndex) Each(ctx context.Context) <-chan restic.PackedBlob {
	mi.idxMutex.RLock()

	ch := make(chan restic.PackedBlob)

	go func() {
		defer mi.idxMutex.RUnlock()
		defer func() {
			close(ch)
		}()

		for _, idx := range mi.idx {
			idxCh := idx.Each(ctx)
			for pb := range idxCh {
				select {
				case <-ctx.Done():
					return
				case ch <- pb:
				}
			}
		}
	}()

	return ch
}

// MergeFinalIndexes merges all final indexes together.
// After calling, there will be only one big final index in MasterIndex
// containing all final index contents.
// Indexes that are not final are left untouched.
// This merging can only be called after all index files are loaded - as
// removing of superseded index contents is only possible for unmerged indexes.
func (mi *MasterIndex) MergeFinalIndexes() {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	// The first index is always final and the one to merge into
	newIdx := mi.idx[:1]
	for i := 1; i < len(mi.idx); i++ {
		idx := mi.idx[i]
		// clear reference in masterindex as it may become stale
		mi.idx[i] = nil
		if !idx.Final() {
			newIdx = append(newIdx, idx)
		} else {
			mi.idx[0].merge(idx)
		}
	}
	mi.idx = newIdx
}

const saveIndexParallelism = 4

// Save saves all known indexes to index files, leaving out any
// packs whose ID is contained in packBlacklist from finalized indexes.
// The new index contains the IDs of all known indexes in the "supersedes"
// field. The IDs are also returned in the IDSet obsolete.
// After calling this function, you should remove the obsolete index files.
func (mi *MasterIndex) Save(ctx context.Context, repo restic.Repository, packBlacklist restic.IDSet, extraObsolete restic.IDs, p *progress.Counter) (obsolete restic.IDSet, err error) {
	mi.idxMutex.Lock()
	defer mi.idxMutex.Unlock()

	debug.Log("start rebuilding index of %d indexes, pack blacklist: %v", len(mi.idx), packBlacklist)

	newIndex := NewIndex()
	obsolete = restic.NewIDSet()

	// track spawned goroutines using wg, create a new context which is
	// cancelled as soon as an error occurs.
	wg, ctx := errgroup.WithContext(ctx)

	ch := make(chan *Index)

	wg.Go(func() error {
		defer close(ch)
		for i, idx := range mi.idx {
			if idx.Final() {
				ids, err := idx.IDs()
				if err != nil {
					debug.Log("index %d does not have an ID: %v", err)
					return err
				}

				debug.Log("adding index ids %v to supersedes field", ids)

				err = newIndex.AddToSupersedes(ids...)
				if err != nil {
					return err
				}
				obsolete.Merge(restic.NewIDSet(ids...))
			} else {
				debug.Log("index %d isn't final, don't add to supersedes field", i)
			}

			debug.Log("adding index %d", i)

			for pbs := range idx.EachByPack(ctx, packBlacklist) {
				newIndex.StorePack(pbs.packID, pbs.blobs)
				p.Add(1)
				if IndexFull(newIndex) {
					select {
					case ch <- newIndex:
					case <-ctx.Done():
						return nil
					}
					newIndex = NewIndex()
				}
			}
		}

		err = newIndex.AddToSupersedes(extraObsolete...)
		if err != nil {
			return err
		}
		obsolete.Merge(restic.NewIDSet(extraObsolete...))

		select {
		case ch <- newIndex:
		case <-ctx.Done():
		}
		return nil
	})

	// a worker receives an index from ch, and saves the index
	worker := func() error {
		for idx := range ch {
			idx.Finalize()
			if _, err := SaveIndex(ctx, repo, idx); err != nil {
				return err
			}
		}
		return nil
	}

	// run workers on ch
	wg.Go(func() error {
		return RunWorkers(saveIndexParallelism, worker)
	})

	err = wg.Wait()

	return obsolete, err
}
<|endoftext|>


<commit_before>// Package simulator provides an interface with package consensus without I/O.
package simulator

import (
	"github.com/heidi-ann/ios/consensus"
	"github.com/heidi-ann/ios/msgs"
)

func RunSimulator(nodes int) []*msgs.Io {
	ios := make([]*msgs.Io, nodes)

	// setup state
	for id := 0; id < nodes; id++ {
		io := msgs.MakeIo(10, nodes)
		conf := consensus.Config{ID: id, N: nodes, LogLength: 1000}
		go consensus.Init(io, conf)
		go io.DumpPersistentStorage()
		ios[id] = io
	}

	// forward traffic
	for to := range ios {
		for from := range ios {
			go ios[to].Incoming.Forward(ios[from].OutgoingUnicast[to])
		}
	}

	return ios
}

// same as RunSimulator except where the log in persistent storage is given
func RunRecoverySimulator(nodes int, logs [][]msgs.Entry, views []int) []*msgs.Io {
	ios := make([]*msgs.Io, nodes)

	// setup state
	for id := 0; id < nodes; id++ {
		io := msgs.MakeIo(10, nodes)
		conf := consensus.Config{ID: id, N: nodes, LogLength: 1000}
		go consensus.Recover(io, conf, views[id], logs[id])
		go io.DumpPersistentStorage()
		ios[id] = io
	}

	// forward traffic
	for to := range ios {
		for from := range ios {
			go ios[to].Incoming.Forward(ios[from].OutgoingUnicast[to])
		}
	}

	return ios
}
<commit_msg>fixing travis build<commit_after>// Package simulator provides an interface with package consensus without I/O.
package simulator

import (
	"github.com/heidi-ann/ios/consensus"
	"github.com/heidi-ann/ios/msgs"
	"github.com/heidi-ann/ios/store"
)

func RunSimulator(nodes int) []*msgs.Io {
	ios := make([]*msgs.Io, nodes)
  store := store.New()
	// setup state
	for id := 0; id < nodes; id++ {
		io := msgs.MakeIo(10, nodes)
		conf := consensus.Config{ID: id, N: nodes, LogLength: 1000}
		go consensus.Init(io, conf, store)
		go io.DumpPersistentStorage()
		ios[id] = io
	}

	// forward traffic
	for to := range ios {
		for from := range ios {
			go ios[to].Incoming.Forward(ios[from].OutgoingUnicast[to])
		}
	}

	return ios
}

// same as RunSimulator except where the log in persistent storage is given
func RunRecoverySimulator(nodes int, logs [][]msgs.Entry, views []int) []*msgs.Io {
	ios := make([]*msgs.Io, nodes)
	store := store.New()
	// setup state
	for id := 0; id < nodes; id++ {
		io := msgs.MakeIo(10, nodes)
		conf := consensus.Config{ID: id, N: nodes, LogLength: 1000}
		go consensus.Recover(io, conf, views[id], logs[id], store, -1)
		go io.DumpPersistentStorage()
		ios[id] = io
	}

	// forward traffic
	for to := range ios {
		for from := range ios {
			go ios[to].Incoming.Forward(ios[from].OutgoingUnicast[to])
		}
	}

	return ios
}
<|endoftext|>


<commit_before>// Copyright 2009 Dimiter Stanev, malkia@gmail.com. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"os"
	"fmt"
	"io/ioutil"
	"go/parser"
	"go/ast"
	"strconv"
	"path"
)

func getmap(m map[string]string, k string) (v string) {
	v, ok := m[k]
	if !ok {
		v = ""
	}
	return
}

var (
	curdir, _ = os.Getwd()
	envbin    = os.Getenv("GOBIN")
	arch      = getmap(map[string]string{"amd64": "6", "386": "8", "arm": "5"}, os.Getenv("GOARCH"))
)

func exec(args []string, dir string) {
	p, error := os.ForkExec(args[0], args, os.Environ(), dir, []*os.File{os.Stdin, os.Stdout, os.Stderr})
	if error != nil {
		fmt.Fprintf( os.Stderr, "Can't %s\n", error );
		os.Exit(1);
	}
	m, error := os.Wait(p, 0)
	if error != nil {
		fmt.Fprintf( os.Stderr, "Can't %s\n", error );
		os.Exit(1);
	}
	if m.WaitStatus != 0 {
		os.Exit(int(m.WaitStatus));
	}
}

func getLocalImports(filename string) (imports map[string]bool, error os.Error) {
	source, error := ioutil.ReadFile(filename)
	if error != nil {
		return
	}
	file, error := parser.ParseFile(filename, source, parser.ImportsOnly)
	if error != nil {
		return
	}
	for _, importDecl := range file.Decls {
		importDecl, ok := importDecl.(*ast.GenDecl)
		if ok {
			for _, importSpec := range importDecl.Specs {
				importSpec, ok := importSpec.(*ast.ImportSpec)
				if ok {
					for _, importPath := range importSpec.Path {
						importPath, _ := strconv.Unquote(string(importPath.Value))
						if len(importPath) > 0 && importPath[0] == '.' {
							if imports == nil {
								imports = make(map[string]bool)
							}
							dir, _ := path.Split(filename)
							imports[path.Join(dir, path.Clean(importPath))] = true
						}
					}
				}
			}
		}
	}

	return
}

func collectSourceFiles(sourcePath string, sourceTable map[int]string, sourceSet map[string]int) (error os.Error) {
	sourcePath = path.Clean(sourcePath)

	if index, exists := sourceSet[sourcePath]; exists {
		sourceTable[index] = ""
		sourceSet[sourcePath] = 0, false
	}

	localImports, error := getLocalImports(sourcePath + ".go")
	if error != nil {
		return
	}

	index := len(sourceTable)
	sourceSet[sourcePath] = index
	sourceTable[index] = sourcePath

	for k, _ := range localImports {
		if error = collectSourceFiles(k, sourceTable, sourceSet); error != nil {
			break
		}
	}

	return
}

func CollectSourceFiles(sourcePath string) (sourceTable map[int]string, error os.Error) {
	sourceTable = make(map[int]string)
	return sourceTable, collectSourceFiles(sourcePath, sourceTable, make(map[string]int))
}

func shouldUpdate(sourceFile, targetFile string) (doUpdate bool, error os.Error) {
	sourceStat, error := os.Lstat(sourceFile)
	if error != nil {
		return false, error
	}
	targetStat, error := os.Lstat(targetFile)
	if error != nil {
		return true, error
	}
	return targetStat.Mtime_ns < sourceStat.Mtime_ns, error
}

func compile(target string) {
	dir, filename := path.Split(target)
	dir = path.Join(curdir, dir)
	source := path.Join(dir, filename+".go")
	object := path.Join(dir, filename+"."+arch)
	doUpdate, error := shouldUpdate(source, object)
	if doUpdate {
		exec([]string{path.Join(envbin, arch+"g"), filename + ".go"}, dir)
	} else if error != nil {
		fmt.Fprintln(os.Stderr, error)
	}
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "go main-program [arg0 [arg1 ...]]")
		os.Exit(1)
	}

	target := path.Clean(args[0])
	if path.Ext(target) == ".go" {
		target = target[0 : len(target)-3]
	}

	files, error := CollectSourceFiles(target)
	if error != nil {
		fmt.Fprintf(os.Stderr, "Can't %v\n", error)
		os.Exit(1)
	}

	// Compiling source files
	for k := len(files) - 1; k >= 0; k-- {
		v := files[k]
		if v != "" {
			compile(v)
		} else {
			files[k] = "", false
		}
	}

	targets := make([]string, len(files)+3)
	targets[0] = path.Join(envbin, arch+"l")
	targets[1] = "-o"
	targets[2] = target
	doLink := false
	for i, v := range files {
		targets[i+3] = v + "." + arch
		if !doLink {
			if shouldUpdate, _ := shouldUpdate(targets[i+3], target); shouldUpdate {
				doLink = true
			}
		}
	}
	if doLink {
		exec(targets, "")
	}
	os.Exec(path.Join(curdir, target), args, os.Environ())
	fmt.Fprintf(os.Stderr, "Error running %v\n", args)
	os.Exit(1)
}
<commit_msg>fix code to work with recent linker.<commit_after>// Copyright 2009 Dimiter Stanev, malkia@gmail.com. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"os"
	"fmt"
	"io/ioutil"
	"go/parser"
	"go/ast"
	"strconv"
	"path"
)

func getmap(m map[string]string, k string) (v string) {
	v, ok := m[k]
	if !ok {
		v = ""
	}
	return
}

var (
	curdir, _ = os.Getwd()
	envbin    = os.Getenv("GOBIN")
	arch      = getmap(map[string]string{"amd64": "6", "386": "8", "arm": "5"}, os.Getenv("GOARCH"))
)

func exec(args []string, dir string) {
	p, error := os.ForkExec(args[0], args, os.Environ(), dir, []*os.File{os.Stdin, os.Stdout, os.Stderr})
	if error != nil {
		fmt.Fprintf( os.Stderr, "Can't %s\n", error );
		os.Exit(1);
	}
	m, error := os.Wait(p, 0)
	if error != nil {
		fmt.Fprintf( os.Stderr, "Can't %s\n", error );
		os.Exit(1);
	}
	if m.WaitStatus != 0 {
		os.Exit(int(m.WaitStatus));
	}
}

func getLocalImports(filename string) (imports map[string]bool, error os.Error) {
	source, error := ioutil.ReadFile(filename)
	if error != nil {
		return
	}
	file, error := parser.ParseFile(filename, source, parser.ImportsOnly)
	if error != nil {
		return
	}
	for _, importDecl := range file.Decls {
		importDecl, ok := importDecl.(*ast.GenDecl)
		if ok {
			for _, importSpec := range importDecl.Specs {
				importSpec, ok := importSpec.(*ast.ImportSpec)
				if ok {
					for _, importPath := range importSpec.Path {
						importPath, _ := strconv.Unquote(string(importPath.Value))
						if len(importPath) > 0 && importPath[0] == '.' {
							if imports == nil {
								imports = make(map[string]bool)
							}
							dir, _ := path.Split(filename)
							imports[path.Join(dir, path.Clean(importPath))] = true
						}
					}
				}
			}
		}
	}

	return
}

func collectSourceFiles(sourcePath string, sourceTable map[int]string, sourceSet map[string]int) (error os.Error) {
	sourcePath = path.Clean(sourcePath)

	if index, exists := sourceSet[sourcePath]; exists {
		sourceTable[index] = ""
		sourceSet[sourcePath] = 0, false
	}

	localImports, error := getLocalImports(sourcePath + ".go")
	if error != nil {
		return
	}

	index := len(sourceTable)
	sourceSet[sourcePath] = index
	sourceTable[index] = sourcePath

	for k, _ := range localImports {
		if error = collectSourceFiles(k, sourceTable, sourceSet); error != nil {
			break
		}
	}

	return
}

func CollectSourceFiles(sourcePath string) (sourceTable map[int]string, error os.Error) {
	sourceTable = make(map[int]string)
	return sourceTable, collectSourceFiles(sourcePath, sourceTable, make(map[string]int))
}

func shouldUpdate(sourceFile, targetFile string) (doUpdate bool, error os.Error) {
	sourceStat, error := os.Lstat(sourceFile)
	if error != nil {
		return false, error
	}
	targetStat, error := os.Lstat(targetFile)
	if error != nil {
		return true, error
	}
	return targetStat.Mtime_ns < sourceStat.Mtime_ns, error
}

func compile(target string) {
	dir, filename := path.Split(target)
	dir = path.Join(curdir, dir)
	source := path.Join(dir, filename+".go")
	object := path.Join(dir, filename+"."+arch)
	doUpdate, error := shouldUpdate(source, object)
	if doUpdate {
		exec([]string{path.Join(envbin, arch+"g"), filename + ".go"}, dir)
	} else if error != nil {
		fmt.Fprintln(os.Stderr, error)
	}
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "go main-program [arg0 [arg1 ...]]")
		os.Exit(1)
	}

	target := path.Clean(args[0])
	if path.Ext(target) == ".go" {
		target = target[0 : len(target)-3]
	}

	files, error := CollectSourceFiles(target)
	if error != nil {
		fmt.Fprintf(os.Stderr, "Can't %v\n", error)
		os.Exit(1)
	}

	// Compiling source files
	for k := len(files) - 1; k >= 0; k-- {
		v := files[k]
		if v != "" {
			compile(v)
		} else {
			files[k] = "", false
		}
	}

	doLink, _ := shouldUpdate(target+"."+arch, target)
	if doLink {
		exec([]string{path.Join(envbin, arch+"l"), "-o", target, target+"."+arch}, "")
	}
	os.Exec(path.Join(curdir, target), args, os.Environ())
	fmt.Fprintf(os.Stderr, "Error running %v\n", args)
	os.Exit(1)
}
<|endoftext|>