<commit_before>package backend

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/compose-spec/compose-go/cli"
	"github.com/compose-spec/compose-go/types"
	"github.com/docker/ecs-plugin/pkg/compose"
	"github.com/docker/ecs-plugin/pkg/progress"
)

func (b *Backend) Up(ctx context.Context, options *cli.ProjectOptions) error {
	project, err := cli.ProjectFromOptions(options)
	if err != nil {
		return err
	}

	err = b.api.CheckRequirements(ctx, b.Region)
	if err != nil {
		return err
	}

	cluster, err := b.GetCluster(ctx, project)
	if err != nil {
		return err
	}

	template, err := b.Convert(project)
	if err != nil {
		return err
	}

	vpc, err := b.GetVPC(ctx, project)
	if err != nil {
		return err
	}

	subNets, err := b.api.GetSubNets(ctx, vpc)
	if err != nil {
		return err
	}
	if len(subNets) < 2 {
		return fmt.Errorf("VPC %s should have at least 2 associated subnets in different availability zones", vpc)
	}

	lb, err := b.GetLoadBalancer(ctx, project)
	if err != nil {
		return err
	}

	parameters := map[string]string{
		ParameterClusterName:     cluster,
		ParameterVPCId:           vpc,
		ParameterSubnet1Id:       subNets[0],
		ParameterSubnet2Id:       subNets[1],
		ParameterLoadBalancerARN: lb,
	}

	update, err := b.api.StackExists(ctx, project.Name)
	if err != nil {
		return err
	}
	operation := compose.StackCreate
	if update {
		operation = compose.StackUpdate
		changeset, err := b.api.CreateChangeSet(ctx, project.Name, template, parameters)
		if err != nil {
			return err
		}
		err = b.api.UpdateStack(ctx, changeset)
		if err != nil {
			return err
		}
	} else {
		err = b.api.CreateStack(ctx, project.Name, template, parameters)
		if err != nil {
			return err
		}
	}

	for k := range template.Resources {
		b.writer.Event(progress.Event{
			ID:         k,
			Status:     progress.Working,
			StatusText: "Pending",
		})
	}

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-signalChan
		fmt.Println("user interrupted deployment. Deleting stack...")
		b.Down(ctx, options)
	}()

	err = b.WaitStackCompletion(ctx, project.Name, operation)
	// update status for external resources (LB and cluster)
	loadBalancerName := fmt.Sprintf("%.32s", fmt.Sprintf("%sLoadBalancer", strings.Title(project.Name)))
	for k := range template.Resources {
		switch k {
		case "Cluster":
			if cluster == "" {
				continue
			}
		case loadBalancerName:
			if lb == "" {
				continue
			}
		default:
			continue
		}
		b.writer.Event(progress.Event{
			ID:         k,
			Status:     progress.Done,
			StatusText: "",
		})
	}
	return err
}

func (b Backend) GetVPC(ctx context.Context, project *types.Project) (string, error) {
	//check compose file for custom VPC selected
	if vpc, ok := project.Extensions[compose.ExtensionVPC]; ok {
		vpcID := vpc.(string)
		ok, err := b.api.VpcExists(ctx, vpcID)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", fmt.Errorf("VPC does not exist: %s", vpc)
		}
		return vpcID, nil
	}
	defaultVPC, err := b.api.GetDefaultVPC(ctx)
	if err != nil {
		return "", err
	}
	return defaultVPC, nil
}

func (b Backend) GetLoadBalancer(ctx context.Context, project *types.Project) (string, error) {
	//check compose file for custom VPC selected
	if ext, ok := project.Extensions[compose.ExtensionLB]; ok {
		lb := ext.(string)
		ok, err := b.api.LoadBalancerExists(ctx, lb)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", fmt.Errorf("load balancer does not exist: %s", lb)
		}
		return lb, nil
	}
	return "", nil
}

func (b Backend) GetCluster(ctx context.Context, project *types.Project) (string, error) {
	//check compose file for custom VPC selected
	if ext, ok := project.Extensions[compose.ExtensionCluster]; ok {
		cluster := ext.(string)
		ok, err := b.api.ClusterExists(ctx, cluster)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", fmt.Errorf("cluster does not exist: %s", cluster)
		}
		return cluster, nil
	}
	return "", nil
}
<commit_msg>Don't pretend we know resources to be created<commit_after>package backend

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/compose-spec/compose-go/cli"
	"github.com/compose-spec/compose-go/types"
	"github.com/docker/ecs-plugin/pkg/compose"
)

func (b *Backend) Up(ctx context.Context, options *cli.ProjectOptions) error {
	project, err := cli.ProjectFromOptions(options)
	if err != nil {
		return err
	}

	err = b.api.CheckRequirements(ctx, b.Region)
	if err != nil {
		return err
	}

	cluster, err := b.GetCluster(ctx, project)
	if err != nil {
		return err
	}

	template, err := b.Convert(project)
	if err != nil {
		return err
	}

	vpc, err := b.GetVPC(ctx, project)
	if err != nil {
		return err
	}

	subNets, err := b.api.GetSubNets(ctx, vpc)
	if err != nil {
		return err
	}
	if len(subNets) < 2 {
		return fmt.Errorf("VPC %s should have at least 2 associated subnets in different availability zones", vpc)
	}

	lb, err := b.GetLoadBalancer(ctx, project)
	if err != nil {
		return err
	}

	parameters := map[string]string{
		ParameterClusterName:     cluster,
		ParameterVPCId:           vpc,
		ParameterSubnet1Id:       subNets[0],
		ParameterSubnet2Id:       subNets[1],
		ParameterLoadBalancerARN: lb,
	}

	update, err := b.api.StackExists(ctx, project.Name)
	if err != nil {
		return err
	}
	operation := compose.StackCreate
	if update {
		operation = compose.StackUpdate
		changeset, err := b.api.CreateChangeSet(ctx, project.Name, template, parameters)
		if err != nil {
			return err
		}
		err = b.api.UpdateStack(ctx, changeset)
		if err != nil {
			return err
		}
	} else {
		err = b.api.CreateStack(ctx, project.Name, template, parameters)
		if err != nil {
			return err
		}
	}

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-signalChan
		fmt.Println("user interrupted deployment. Deleting stack...")
		b.Down(ctx, options)
	}()

	err = b.WaitStackCompletion(ctx, project.Name, operation)
	return err
}

func (b Backend) GetVPC(ctx context.Context, project *types.Project) (string, error) {
	//check compose file for custom VPC selected
	if vpc, ok := project.Extensions[compose.ExtensionVPC]; ok {
		vpcID := vpc.(string)
		ok, err := b.api.VpcExists(ctx, vpcID)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", fmt.Errorf("VPC does not exist: %s", vpc)
		}
		return vpcID, nil
	}
	defaultVPC, err := b.api.GetDefaultVPC(ctx)
	if err != nil {
		return "", err
	}
	return defaultVPC, nil
}

func (b Backend) GetLoadBalancer(ctx context.Context, project *types.Project) (string, error) {
	//check compose file for custom VPC selected
	if ext, ok := project.Extensions[compose.ExtensionLB]; ok {
		lb := ext.(string)
		ok, err := b.api.LoadBalancerExists(ctx, lb)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", fmt.Errorf("load balancer does not exist: %s", lb)
		}
		return lb, nil
	}
	return "", nil
}

func (b Backend) GetCluster(ctx context.Context, project *types.Project) (string, error) {
	//check compose file for custom VPC selected
	if ext, ok := project.Extensions[compose.ExtensionCluster]; ok {
		cluster := ext.(string)
		ok, err := b.api.ClusterExists(ctx, cluster)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", fmt.Errorf("cluster does not exist: %s", cluster)
		}
		return cluster, nil
	}
	return "", nil
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"net/http"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/spanner"
	"golang.org/x/sync/errgroup"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"

	"go.chromium.org/luci/common/bq"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/resultdb/internal"
	"go.chromium.org/luci/resultdb/internal/span"
	bqpb "go.chromium.org/luci/resultdb/proto/bq/v1"
	pb "go.chromium.org/luci/resultdb/proto/rpc/v1"
)

const (
	maxInvocationGraphSize = 1000
	maxBatchSize           = 500
)

// inserter is implemented by bigquery.Inserter.
type inserter interface {
	// Put uploads one or more rows to the BigQuery service.
	Put(ctx context.Context, src interface{}) error
}

func getLUCIProject(ctx context.Context, invID span.InvocationID) (string, error) {
	realm, err := span.ReadInvocationRealm(ctx, span.Client(ctx).Single(), invID)
	if err != nil {
		return "", err
	}

	project, _, err := internal.ParseRealm(realm)
	if err != nil {
		return "", errors.Annotate(err, "invocation %q", invID.Name()).Err()
	}
	return project, nil
}

func getBQClient(ctx context.Context, luciProject string, bqExport *pb.BigQueryExport) (*bigquery.Client, error) {
	tr, err := auth.GetRPCTransport(ctx, auth.AsProject, auth.WithProject(luciProject), auth.WithScopes(bigquery.Scope))
	if err != nil {
		return nil, err
	}

	return bigquery.NewClient(ctx, bqExport.Project, option.WithHTTPClient(&http.Client{
		Transport: tr,
	}))
}

// ensureBQTable creates a BQ table if it doesn't exist.
func ensureBQTable(ctx context.Context, client *bigquery.Client, bqExport *pb.BigQueryExport) error {
	t := client.Dataset(bqExport.Dataset).Table(bqExport.Table)

	// Check the existence of table.
	// TODO(chanli): Cache the check result.
	_, err := t.Metadata(ctx)
	apiErr, ok := err.(*googleapi.Error)
	switch {
	case ok && apiErr.Code == http.StatusNotFound:
		// Table doesn't exist.
		break
	case ok && apiErr.Code == http.StatusForbidden:
		// No read table permission.
		return permanentInvocationTaskErrTag.Apply(err)
	default:
		// Either no err or the err is not special cases above, simply return.
		return err
	}

	// Table doesn't exist. Create one.
	err = t.Create(ctx, nil)
	apiErr, ok = err.(*googleapi.Error)
	switch {
	case err == nil:
		logging.Infof(ctx, "Created BigQuery table %s.%s.%s", bqExport.Project, bqExport.Dataset, bqExport.Table)
		return nil
	case ok && apiErr.Code == http.StatusConflict:
		// Table just got created. This is fine.
		return nil
	case ok && apiErr.Code == http.StatusForbidden:
		// No create table permission.
		return permanentInvocationTaskErrTag.Apply(err)
	default:
		return err
	}
}

type testVariantKey struct {
	testID      string
	variantHash string
}

// queryExoneratedTestVariants reads exonerated test variants matching the predicate.
func queryExoneratedTestVariants(ctx context.Context, txn *spanner.ReadOnlyTransaction, invIDs span.InvocationIDSet) (map[testVariantKey]struct{}, error) {
	st := spanner.NewStatement(`
		SELECT DISTINCT TestId, VariantHash,
		FROM TestExonerations
		WHERE InvocationId IN UNNEST(@invIDs)
	`)
	st.Params["invIDs"] = invIDs

	tvs := map[testVariantKey]struct{}{}
	var b span.Buffer
	err := span.Query(ctx, "exonerated test variants", txn, st, func(row *spanner.Row) error {
		var key testVariantKey
		if err := b.FromSpanner(row, &key.testID, &key.variantHash); err != nil {
			return err
		}
		tvs[key] = struct{}{}
		return nil
	})
	if err != nil {
		return nil, err
	}

	return tvs, nil
}

func generateBQRow(inv *pb.Invocation, tr *pb.TestResult, exonerated bool) *bq.Row {
	return &bq.Row{
		Message: &bqpb.TestResultRow{
			Invocation: &bqpb.TestResultRow_Invocation{
				Id:          string(span.MustParseInvocationName(inv.Name)),
				Interrupted: inv.Interrupted,
				Tags:        inv.Tags,
			},
			Result: tr,
			Exoneration: &bqpb.TestResultRow_TestExoneration{
				Exonerated: exonerated,
			},
		},
		InsertID: tr.Name,
	}
}

func queryTestResultsStreaming(ctx context.Context, txn *spanner.ReadOnlyTransaction, inv *pb.Invocation, q span.TestResultQuery, exoneratedTestVariants map[testVariantKey]struct{}, maxBatchSize int, batchC chan []*bq.Row) error {
	rows := make([]*bq.Row, 0, maxBatchSize)
	err := span.QueryTestResultsStreaming(ctx, txn, q, func(tr *pb.TestResult, variantHash string) error {
		_, exonerated := exoneratedTestVariants[testVariantKey{testID: tr.TestId, variantHash: variantHash}]
		rows = append(rows, generateBQRow(inv, tr, exonerated))
		if len(rows) >= maxBatchSize {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case batchC <- rows:
			}
			rows = make([]*bq.Row, 0, maxBatchSize)
		}
		return nil
	})

	if err != nil {
		return err
	}

	if len(rows) > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case batchC <- rows:
		}
	}

	return nil
}

func batchExportRows(ctx context.Context, ins inserter, batchC chan []*bq.Row) error {
	return parallel.WorkPool(10, func(workC chan<- func() error) {
		for rows := range batchC {
			rows := rows
			workC <- func() error {
				err := ins.Put(ctx, rows)
				if apiErr, ok := err.(*googleapi.Error); ok && apiErr.Code == http.StatusForbidden {
					err = permanentInvocationTaskErrTag.Apply(err)
				}
				return err
			}
		}
	})
}

// exportTestResultsToBigQuery queries test results in Spanner then exports them to BigQuery.
func exportTestResultsToBigQuery(ctx context.Context, ins inserter, invID span.InvocationID, bqExport *pb.BigQueryExport, maxBatchSize int) error {
	txn := span.Client(ctx).ReadOnlyTransaction()
	defer txn.Close()

	inv, err := span.ReadInvocationFull(ctx, txn, invID)
	if err != nil {
		return err
	}
	if inv.State != pb.Invocation_FINALIZED {
		return errors.Reason("%s is not finalized yet", invID.Name()).Err()
	}

	// Get the invocation set.
	invIDs, err := span.ReadReachableInvocations(ctx, txn, maxInvocationGraphSize, span.NewInvocationIDSet(invID))
	if err != nil {
		if span.TooManyInvocationsTag.In(err) {
			err = permanentInvocationTaskErrTag.Apply(err)
		}
		return err
	}

	exoneratedTestVariants, err := queryExoneratedTestVariants(ctx, txn, invIDs)
	if err != nil {
		return err
	}

	// Query test results and export to BigQuery.
	batchC := make(chan []*bq.Row)

	// Batch exports rows to BigQuery.
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
		return batchExportRows(ctx, ins, batchC)
	})

	q := span.TestResultQuery{
		Predicate:         bqExport.GetTestResults().GetPredicate(),
		InvocationIDs:     invIDs,
		SelectVariantHash: true,
	}
	eg.Go(func() error {
		defer close(batchC)
		return queryTestResultsStreaming(ctx, txn, inv, q, exoneratedTestVariants, maxBatchSize, batchC)
	})

	return eg.Wait()
}

// exportResultsToBigQuery exports results of an invocation to a BigQuery table.
func exportResultsToBigQuery(ctx context.Context, invID span.InvocationID, bqExport *pb.BigQueryExport) error {
	luciProject, err := getLUCIProject(ctx, invID)
	if err != nil {
		return err
	}

	client, err := getBQClient(ctx, luciProject, bqExport)
	if err != nil {
		return err
	}
	defer client.Close()

	if err := ensureBQTable(ctx, client, bqExport); err != nil {
		return err
	}

	ins := client.Dataset(bqExport.Dataset).Table(bqExport.Table).Inserter()
	return exportTestResultsToBigQuery(ctx, ins, invID, bqExport, maxBatchSize)
}
<commit_msg>[resultdb] Trim artifact fetch_url<commit_after>// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"net/http"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/spanner"
	"golang.org/x/sync/errgroup"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"

	"go.chromium.org/luci/common/bq"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/resultdb/internal"
	"go.chromium.org/luci/resultdb/internal/span"
	bqpb "go.chromium.org/luci/resultdb/proto/bq/v1"
	pb "go.chromium.org/luci/resultdb/proto/rpc/v1"
)

const (
	maxInvocationGraphSize = 1000
	maxBatchSize           = 500
)

// inserter is implemented by bigquery.Inserter.
type inserter interface {
	// Put uploads one or more rows to the BigQuery service.
	Put(ctx context.Context, src interface{}) error
}

func getLUCIProject(ctx context.Context, invID span.InvocationID) (string, error) {
	realm, err := span.ReadInvocationRealm(ctx, span.Client(ctx).Single(), invID)
	if err != nil {
		return "", err
	}

	project, _, err := internal.ParseRealm(realm)
	if err != nil {
		return "", errors.Annotate(err, "invocation %q", invID.Name()).Err()
	}
	return project, nil
}

func getBQClient(ctx context.Context, luciProject string, bqExport *pb.BigQueryExport) (*bigquery.Client, error) {
	tr, err := auth.GetRPCTransport(ctx, auth.AsProject, auth.WithProject(luciProject), auth.WithScopes(bigquery.Scope))
	if err != nil {
		return nil, err
	}

	return bigquery.NewClient(ctx, bqExport.Project, option.WithHTTPClient(&http.Client{
		Transport: tr,
	}))
}

// ensureBQTable creates a BQ table if it doesn't exist.
func ensureBQTable(ctx context.Context, client *bigquery.Client, bqExport *pb.BigQueryExport) error {
	t := client.Dataset(bqExport.Dataset).Table(bqExport.Table)

	// Check the existence of table.
	// TODO(chanli): Cache the check result.
	_, err := t.Metadata(ctx)
	apiErr, ok := err.(*googleapi.Error)
	switch {
	case ok && apiErr.Code == http.StatusNotFound:
		// Table doesn't exist.
		break
	case ok && apiErr.Code == http.StatusForbidden:
		// No read table permission.
		return permanentInvocationTaskErrTag.Apply(err)
	default:
		// Either no err or the err is not special cases above, simply return.
		return err
	}

	// Table doesn't exist. Create one.
	err = t.Create(ctx, nil)
	apiErr, ok = err.(*googleapi.Error)
	switch {
	case err == nil:
		logging.Infof(ctx, "Created BigQuery table %s.%s.%s", bqExport.Project, bqExport.Dataset, bqExport.Table)
		return nil
	case ok && apiErr.Code == http.StatusConflict:
		// Table just got created. This is fine.
		return nil
	case ok && apiErr.Code == http.StatusForbidden:
		// No create table permission.
		return permanentInvocationTaskErrTag.Apply(err)
	default:
		return err
	}
}

type testVariantKey struct {
	testID      string
	variantHash string
}

// queryExoneratedTestVariants reads exonerated test variants matching the predicate.
func queryExoneratedTestVariants(ctx context.Context, txn *spanner.ReadOnlyTransaction, invIDs span.InvocationIDSet) (map[testVariantKey]struct{}, error) {
	st := spanner.NewStatement(`
		SELECT DISTINCT TestId, VariantHash,
		FROM TestExonerations
		WHERE InvocationId IN UNNEST(@invIDs)
	`)
	st.Params["invIDs"] = invIDs

	tvs := map[testVariantKey]struct{}{}
	var b span.Buffer
	err := span.Query(ctx, "exonerated test variants", txn, st, func(row *spanner.Row) error {
		var key testVariantKey
		if err := b.FromSpanner(row, &key.testID, &key.variantHash); err != nil {
			return err
		}
		tvs[key] = struct{}{}
		return nil
	})
	if err != nil {
		return nil, err
	}

	return tvs, nil
}

func generateBQRow(inv *pb.Invocation, tr *pb.TestResult, exonerated bool) *bq.Row {
	return &bq.Row{
		Message: &bqpb.TestResultRow{
			Invocation: &bqpb.TestResultRow_Invocation{
				Id:          string(span.MustParseInvocationName(inv.Name)),
				Interrupted: inv.Interrupted,
				Tags:        inv.Tags,
			},
			Result: tr,
			Exoneration: &bqpb.TestResultRow_TestExoneration{
				Exonerated: exonerated,
			},
		},
		InsertID: tr.Name,
	}
}

func queryTestResultsStreaming(ctx context.Context, txn *spanner.ReadOnlyTransaction, inv *pb.Invocation, q span.TestResultQuery, exoneratedTestVariants map[testVariantKey]struct{}, maxBatchSize int, batchC chan []*bq.Row) error {
	rows := make([]*bq.Row, 0, maxBatchSize)
	err := span.QueryTestResultsStreaming(ctx, txn, q, func(tr *pb.TestResult, variantHash string) error {
		trimTestResultForBigQuery(tr)

		_, exonerated := exoneratedTestVariants[testVariantKey{testID: tr.TestId, variantHash: variantHash}]
		rows = append(rows, generateBQRow(inv, tr, exonerated))
		if len(rows) >= maxBatchSize {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case batchC <- rows:
			}
			rows = make([]*bq.Row, 0, maxBatchSize)
		}
		return nil
	})

	if err != nil {
		return err
	}

	if len(rows) > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case batchC <- rows:
		}
	}

	return nil
}

func batchExportRows(ctx context.Context, ins inserter, batchC chan []*bq.Row) error {
	return parallel.WorkPool(10, func(workC chan<- func() error) {
		for rows := range batchC {
			rows := rows
			workC <- func() error {
				err := ins.Put(ctx, rows)
				if apiErr, ok := err.(*googleapi.Error); ok && apiErr.Code == http.StatusForbidden {
					err = permanentInvocationTaskErrTag.Apply(err)
				}
				return err
			}
		}
	})
}

// exportTestResultsToBigQuery queries test results in Spanner then exports them to BigQuery.
func exportTestResultsToBigQuery(ctx context.Context, ins inserter, invID span.InvocationID, bqExport *pb.BigQueryExport, maxBatchSize int) error {
	txn := span.Client(ctx).ReadOnlyTransaction()
	defer txn.Close()

	inv, err := span.ReadInvocationFull(ctx, txn, invID)
	if err != nil {
		return err
	}
	if inv.State != pb.Invocation_FINALIZED {
		return errors.Reason("%s is not finalized yet", invID.Name()).Err()
	}

	// Get the invocation set.
	invIDs, err := span.ReadReachableInvocations(ctx, txn, maxInvocationGraphSize, span.NewInvocationIDSet(invID))
	if err != nil {
		if span.TooManyInvocationsTag.In(err) {
			err = permanentInvocationTaskErrTag.Apply(err)
		}
		return err
	}

	exoneratedTestVariants, err := queryExoneratedTestVariants(ctx, txn, invIDs)
	if err != nil {
		return err
	}

	// Query test results and export to BigQuery.
	batchC := make(chan []*bq.Row)

	// Batch exports rows to BigQuery.
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
		return batchExportRows(ctx, ins, batchC)
	})

	q := span.TestResultQuery{
		Predicate:         bqExport.GetTestResults().GetPredicate(),
		InvocationIDs:     invIDs,
		SelectVariantHash: true,
	}
	eg.Go(func() error {
		defer close(batchC)
		return queryTestResultsStreaming(ctx, txn, inv, q, exoneratedTestVariants, maxBatchSize, batchC)
	})

	return eg.Wait()
}

// exportResultsToBigQuery exports results of an invocation to a BigQuery table.
func exportResultsToBigQuery(ctx context.Context, invID span.InvocationID, bqExport *pb.BigQueryExport) error {
	luciProject, err := getLUCIProject(ctx, invID)
	if err != nil {
		return err
	}

	client, err := getBQClient(ctx, luciProject, bqExport)
	if err != nil {
		return err
	}
	defer client.Close()

	if err := ensureBQTable(ctx, client, bqExport); err != nil {
		return err
	}

	ins := client.Dataset(bqExport.Dataset).Table(bqExport.Table).Inserter()
	return exportTestResultsToBigQuery(ctx, ins, invID, bqExport, maxBatchSize)
}

// trimTestResultForBigQuery trims fields that should not go to BigQuery.
func trimTestResultForBigQuery(tr *pb.TestResult) {
	trimArts := func(arts []*pb.Artifact) {
		for _, a := range arts {
			a.FetchUrl = ""
		}
	}
	trimArts(tr.InputArtifacts)
	trimArts(tr.OutputArtifacts)
}
<|endoftext|>


<commit_before>package cli_test

import "testing"

func TestRunDetached(t *testing.T) {
	run(t, []Command{
		DeployCommand("latest", "v1"),
		{
			"run -d migration -a acme-inc",
			"Ran `migration` on acme-inc as run, detached.",
		},
	})
}

func TestRunAttached(t *testing.T) {
	run(t, []Command{
		DeployCommand("latest", "v1"),
		{
			"run migration -a acme-inc",
			"Fake output for `migration` on acme-inc",
		},
	})
}
<commit_msg>Disable run tests.<commit_after>package cli_test

import "testing"

func testRunDetached(t *testing.T) {
	run(t, []Command{
		DeployCommand("latest", "v1"),
		{
			"run -d migration -a acme-inc",
			"Ran `migration` on acme-inc as run, detached.",
		},
	})
}

func testRunAttached(t *testing.T) {
	run(t, []Command{
		DeployCommand("latest", "v1"),
		{
			"run migration -a acme-inc",
			"Fake output for `migration` on acme-inc",
		},
	})
}
<|endoftext|>


<commit_before>package web

import (
	"github.com/gin-gonic/contrib/gzip"
	"github.com/gin-gonic/contrib/static"
	"github.com/gin-gonic/gin"
	"monitor"
	"time"
)

var active_monitor *monitor.Monitor
var start_time time.Time

func New(m *monitor.Monitor) *gin.Engine {
	_ = gzip.Gzip

	active_monitor = m

	gin.SetMode(gin.ReleaseMode)

	r := gin.Default()
	r.Use(gzip.Gzip(gzip.DefaultCompression))
	r.LoadHTMLGlob("src/web/templates/*")

	r.GET("/", getIndex)
	r.GET("/api/snapshots", getSnapshots) //, gzip.Gzip(gzip.DefaultCompression))

	r.Use(static.Serve("/", static.LocalFile("src/web/static/", false)))

	start_time = time.Now()

	return r
}

func getIndex(c *gin.Context) {
	c.HTML(200, "index.html", gin.H{
		"StartTime": start_time,
		"Interval":  int64(active_monitor.Interval / 1e6),
		"Mode":      gin.Mode(),
	})
}

func getSnapshots(c *gin.Context) {
	now := time.Now()

	recent := active_monitor.GetRecentSnapshots().Filter(func(s monitor.Snapshot) bool {
		res := isTimestampInLast(s.Timestamp, now, 60*time.Second) ||
			(isTimestampInLast(s.Timestamp, now, 5*time.Minute) && isSignificantTimestamp(s.Timestamp, 10*time.Second)) ||
			(isTimestampInLast(s.Timestamp, now, 2*time.Hour) && isSignificantTimestamp(s.Timestamp, 10*time.Minute)) ||
			(isTimestampInLast(s.Timestamp, now, 24*time.Hour) && isSignificantTimestamp(s.Timestamp, 30*time.Minute))

		return res
	})

	c.JSON(200, recent)
}

func isTimestampInLast(s, now time.Time, dur time.Duration) bool {
	return now.Sub(s) < dur
}

func isSignificantTimestamp(s time.Time, frequency time.Duration) bool {
	return s.UnixNano()%int64(frequency) < int64(active_monitor.Interval)
}
<commit_msg>Changing frequencies again<commit_after>package web

import (
	"github.com/gin-gonic/contrib/gzip"
	"github.com/gin-gonic/contrib/static"
	"github.com/gin-gonic/gin"
	"monitor"
	"time"
)

var active_monitor *monitor.Monitor
var start_time time.Time

func New(m *monitor.Monitor) *gin.Engine {
	_ = gzip.Gzip

	active_monitor = m

	gin.SetMode(gin.ReleaseMode)

	r := gin.Default()
	r.Use(gzip.Gzip(gzip.DefaultCompression))
	r.LoadHTMLGlob("src/web/templates/*")

	r.GET("/", getIndex)
	r.GET("/api/snapshots", getSnapshots) //, gzip.Gzip(gzip.DefaultCompression))

	r.Use(static.Serve("/", static.LocalFile("src/web/static/", false)))

	start_time = time.Now()

	return r
}

func getIndex(c *gin.Context) {
	c.HTML(200, "index.html", gin.H{
		"StartTime": start_time,
		"Interval":  int64(active_monitor.Interval / 1e6),
		"Mode":      gin.Mode(),
	})
}

func getSnapshots(c *gin.Context) {
	now := time.Now()

	recent := active_monitor.GetRecentSnapshots().Filter(func(s monitor.Snapshot) bool {
		res := isTimestampInLast(s.Timestamp, now, 60*time.Second) ||
			(isTimestampInLast(s.Timestamp, now, 5*time.Minute) && isSignificantTimestamp(s.Timestamp, 10*time.Second)) ||
			(isTimestampInLast(s.Timestamp, now, 2*time.Hour) && isSignificantTimestamp(s.Timestamp, 5*time.Minute)) ||
			(isTimestampInLast(s.Timestamp, now, 48*time.Hour) && isSignificantTimestamp(s.Timestamp, 30*time.Minute))

		return res
	})

	c.JSON(200, recent)
}

func isTimestampInLast(s, now time.Time, dur time.Duration) bool {
	return now.Sub(s) < dur
}

func isSignificantTimestamp(s time.Time, frequency time.Duration) bool {
	return s.UnixNano()%int64(frequency) < int64(active_monitor.Interval)
}
<|endoftext|>


<commit_before>// Паттерн Мост (Bridge)
//

package bridge

// Тип Carer, описывает интерфейс автомобиля
type Carer interface {
	Rase() string // автомобиль умеет ездить (для этого и нужен двигатель)
}

// Тип Enginer, описывает интерфейс двигателя
// Каждый двигатель должен его реализовать
type Enginer interface {
	GetSound() string // Возвращает звук двигателя
}

// Тип Car, реализует автомобиль
type Car struct {
	engine Enginer
}

// Машина едет
func (self *Car) Rase() string {
	return self.engine.GetSound()
}

// Тип EngineSuzuki, реализует двигатель Suzuki
type EngineSuzuki struct {
}

// Метод отвечает за завук двигателя
func (self *EngineSuzuki) GetSound() string {
	return "SssuuuuZzzuuuuKkiiiii"
}

// Тип EngineHonda, реализует двигатель Honda
type EngineHonda struct {
}

// Метод отвечает за завук двигателя
func (self *EngineHonda) GetSound() string {
	return "HhoooNnnnnnnnnDddaaaaaaa"
}

// Тип EngineLada, реализует двигатель АвтоВаза
type EngineLada struct {
}

// Метод отвечает за завук двигателя
func (self *EngineLada) GetSound() string {
	return "PhhhhPhhhhPhPhPhPhPh"
}
<commit_msg>Fixed typo<commit_after>// Паттерн Мост (Bridge)
//

package bridge

// Тип Carer, описывает интерфейс автомобиля
type Carer interface {
	Rase() string // автомобиль умеет ездить (для этого и нужен двигатель)
}

// Тип Enginer, описывает интерфейс двигателя
// Каждый двигатель должен его реализовать
type Enginer interface {
	GetSound() string // Возвращает звук двигателя
}

// Тип Car, реализует автомобиль
type Car struct {
	engine Enginer
}

// Машина едет
func (self *Car) Rase() string {
	return self.engine.GetSound()
}

// Тип EngineSuzuki, реализует двигатель Suzuki
type EngineSuzuki struct {
}

// Метод отвечает за звук двигателя
func (self *EngineSuzuki) GetSound() string {
	return "SssuuuuZzzuuuuKkiiiii"
}

// Тип EngineHonda, реализует двигатель Honda
type EngineHonda struct {
}

// Метод отвечает за звук двигателя
func (self *EngineHonda) GetSound() string {
	return "HhoooNnnnnnnnnDddaaaaaaa"
}

// Тип EngineLada, реализует двигатель АвтоВаза
type EngineLada struct {
}

// Метод отвечает за звук двигателя
func (self *EngineLada) GetSound() string {
	return "PhhhhPhhhhPhPhPhPhPh"
}
<|endoftext|>


<commit_before>// Copyright 2016 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 misc

import (
	"bytes"
	"fmt"

	"github.com/coreos/mantle/kola/register"
	"github.com/coreos/mantle/platform"
)

func init() {
	register.Register(&register.Test{
		Run:         VerityVerify,
		ClusterSize: 1,
		Name:        "coreos.verity.verify",
		Platforms:   []string{"qemu", "aws", "gce"},
	})
}

// VerityVerify asserts that the filesystem mounted on /usr matches the
// dm-verity hash that is embedded in the CoreOS kernel.
func VerityVerify(c platform.TestCluster) error {
	m := c.Machines()[0]

	// extract verity hash from kernel
	hash, err := m.SSH("dd if=/boot/coreos/vmlinuz-a skip=64 count=64 bs=1 2>/dev/null")
	if err != nil {
		return fmt.Errorf("failed to extract verity hash from kernel: %v: %v", hash, err)
	}

	// find /usr dev
	usrdev, err := m.SSH("findmnt -no SOURCE /usr")
	if err != nil {
		return fmt.Errorf("failed to find device for /usr: %v: %v", usrdev, err)
	}

	// figure out partition size for hash dev offset
	offset, err := m.SSH("sudo e2size " + string(usrdev))
	if err != nil {
		return fmt.Errorf("failed to find /usr partition size: %v: %v", offset, err)
	}

	offset = bytes.TrimSpace(offset)
	veritycmd := fmt.Sprintf("sudo veritysetup verify --verbose --hash-offset=%s %s %s %s", offset, usrdev, usrdev, hash)

	verify, err := m.SSH(veritycmd)
	if err != nil {
		return fmt.Errorf("verity hash verification on %s failed: %v: %v", usrdev, verify, err)
	}

	return nil
}
<commit_msg>kola/tests/misc: implement coreos.verity.corruption<commit_after>// Copyright 2016 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 misc

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/coreos/mantle/kola/register"
	"github.com/coreos/mantle/platform"
)

func init() {
	register.Register(&register.Test{
		Run:         VerityVerify,
		ClusterSize: 1,
		Name:        "coreos.verity.verify",
		Platforms:   []string{"qemu", "aws", "gce"},
	})
	register.Register(&register.Test{
		Run:         VerityCorruption,
		ClusterSize: 1,
		Name:        "coreos.verity.corruption",
		Platforms:   []string{"qemu", "aws", "gce"},
	})
}

// Verity verification tests.
// TODO(mischief): seems like a good candidate for kolet.

// VerityVerify asserts that the filesystem mounted on /usr matches the
// dm-verity hash that is embedded in the CoreOS kernel.
func VerityVerify(c platform.TestCluster) error {
	m := c.Machines()[0]

	// extract verity hash from kernel
	hash, err := m.SSH("dd if=/boot/coreos/vmlinuz-a skip=64 count=64 bs=1 2>/dev/null")
	if err != nil {
		return fmt.Errorf("failed to extract verity hash from kernel: %v: %v", hash, err)
	}

	// find /usr dev
	usrdev, err := m.SSH("findmnt -no SOURCE /usr")
	if err != nil {
		return fmt.Errorf("failed to find device for /usr: %v: %v", usrdev, err)
	}

	// XXX: if the /usr dev is /dev/mapper/usr, we're on a verity enabled
	// image, so use dmsetup to find the real device.
	if strings.TrimSpace(string(usrdev)) == "/dev/mapper/usr" {
		usrdev, err = m.SSH("echo -n /dev/$(sudo dmsetup info --noheadings -Co blkdevs_used usr)")
		if err != nil {
			return fmt.Errorf("failed to find device for /usr: %v: %v", usrdev, err)
		}
	}

	// figure out partition size for hash dev offset
	offset, err := m.SSH("sudo e2size " + string(usrdev))
	if err != nil {
		return fmt.Errorf("failed to find /usr partition size: %v: %v", offset, err)
	}

	offset = bytes.TrimSpace(offset)
	veritycmd := fmt.Sprintf("sudo veritysetup verify --verbose --hash-offset=%s %s %s %s", offset, usrdev, usrdev, hash)

	verify, err := m.SSH(veritycmd)
	if err != nil {
		return fmt.Errorf("verity hash verification on %s failed: %v: %v", usrdev, verify, err)
	}

	return nil
}

// VerityCorruption asserts that a machine will fail to read a file from a
// verify filesystem whose blocks have been modified.
func VerityCorruption(c platform.TestCluster) error {
	m := c.Machines()[0]
	// figure out if we are actually using verity
	out, err := m.SSH("sudo veritysetup status usr")
	if err != nil && bytes.Equal(out, []byte("/dev/mapper/usr is inactive.")) {
		// verity not in use, so skip.
		return register.Skip
	} else if err != nil {
		return fmt.Errorf("failed checking verity status: %s: %v", out, err)
	}

	// assert that dm shows verity is in use and the device is valid (V)
	out, err = m.SSH("sudo dmsetup --target verity status usr")
	if err != nil {
		return fmt.Errorf("failed checking dmsetup status of usr: %s: %v", out, err)
	}

	fields := strings.Fields(string(out))
	if len(fields) != 4 {
		return fmt.Errorf("failed checking dmsetup status of usr: not enough fields in output (got %d)", len(fields))
	}

	if fields[3] != "V" {
		return fmt.Errorf("dmsetup status usr reports verity is not valid!")
	}

	// corrupt a file on disk and flush disk caches.
	// try setting NAME=CoreOS to NAME=LulzOS in /usr/lib/os-release

	// get usr device, probably vda3
	usrdev, err := m.SSH("echo /dev/$(sudo dmsetup info --noheadings -Co blkdevs_used usr)")
	if err != nil {
		return fmt.Errorf("failed getting /usr device from dmsetup: %s: %v", out, err)
	}

	// poke bytes into /usr/lib/os-release
	out, err = m.SSH(fmt.Sprintf(`echo NAME=LulzOS | sudo dd of=%s seek=$(expr $(sudo debugfs -R "blocks /lib/os-release" %s 2>/dev/null) \* 4096) bs=1 2>/dev/null`, usrdev, usrdev))
	if err != nil {
		return fmt.Errorf("failed overwriting disk block: %s: %v", out, err)
	}

	// make sure we flush everything so cat has to go through to the device backing verity.
	out, err = m.SSH("sudo /bin/sh -c 'sync; echo -n 3 >/proc/sys/vm/drop_caches'")
	if err != nil {
		return fmt.Errorf("failed dropping disk caches: %s: %v", out, err)
	}

	// read the file back. if we can read it successfully, verity did not do its job.
	out, err = m.SSH("cat /usr/lib/os-release")
	if err == nil {
		return fmt.Errorf("verity did not prevent reading a corrupted file!")
	}

	// assert that dm shows verity device is now corrupted (C)
	out, err = m.SSH("sudo dmsetup --target verity status usr")
	if err != nil {
		return fmt.Errorf("failed checking dmsetup status of usr: %s: %v", out, err)
	}

	fields = strings.Fields(string(out))
	if len(fields) != 4 {
		return fmt.Errorf("failed checking dmsetup status of usr: not enough fields in output (got %d)", len(fields))
	}

	if fields[3] != "C" {
		return fmt.Errorf("dmsetup status usr reports verity is valid after corruption!")
	}

	return nil
}
<|endoftext|>


<commit_before>package bsdiff

import (
	"fmt"
	"os"
	"runtime"
	"sync/atomic"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/itchio/wharf/state"
)

type chunk struct {
	addOldStart int
	addNewStart int
	addLength   int
	copyStart   int
	copyEnd     int
	offset      int
	eoc         bool
}

type SendChunkFunc func(c chunk)

type blockWorkerState struct {
	consumed chan bool
	work     chan int
	chunks   chan chunk
}

func (ctx *DiffContext) doPartitioned(obuf []byte, obuflen int, nbuf []byte, nbuflen int, memstats *runtime.MemStats, writeMessage WriteMessageFunc, consumer *state.Consumer) error {
	var err error

	partitions := ctx.Partitions
	if partitions >= len(obuf)-1 {
		partitions = 1
	}

	consumer.ProgressLabel(fmt.Sprintf("Sorting %s...", humanize.IBytes(uint64(obuflen))))
	consumer.Progress(0.0)

	startTime := time.Now()

	pmemstats := &runtime.MemStats{}
	runtime.ReadMemStats(pmemstats)
	oldAlloc := pmemstats.TotalAlloc

	if ctx.I == nil {
		ctx.I = make([]int, len(obuf))
		beforeAlloc := time.Now()
		fmt.Fprintf(os.Stderr, "\nAllocated %d-int I in %s\n", len(obuf), time.Since(beforeAlloc))
	} else {
		for len(ctx.I) < len(obuf) {
			lenBefore := len(ctx.I)
			beforeAlloc := time.Now()
			ctx.I = make([]int, len(obuf))
			fmt.Fprintf(os.Stderr, "\nGrown I from %d to %d in %s\n", lenBefore, len(ctx.I), time.Since(beforeAlloc))
		}
	}

	psa := NewPSA(partitions, obuf, ctx.I)

	runtime.ReadMemStats(pmemstats)
	newAlloc := pmemstats.TotalAlloc
	fmt.Fprintf(os.Stderr, "\nAlloc difference after PSA: %s. Size of I: %s\n", humanize.IBytes(uint64(newAlloc-oldAlloc)), humanize.IBytes(uint64(8*len(psa.I))))

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentSorting += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after qsufsort: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	bsdc := &Control{}

	consumer.ProgressLabel(fmt.Sprintf("Preparing to scan %s...", humanize.IBytes(uint64(nbuflen))))
	consumer.Progress(0.0)

	startTime = time.Now()

	var waitConsumeTime int64
	var waitEnqueueTime int64
	var idleWorkerTime int64
	var sendingWorkerTime int64

	analyzeBlock := func(nbuflen int, nbuf []byte, offset int, chunks chan chunk) {
		var lenf int

		// Compute the differences, writing ctrl as we go
		var scan, pos, length int
		var lastscan, lastpos, lastoffset int

		for scan < nbuflen {
			var oldscore int
			scan += length

			for scsc := scan; scan < nbuflen; scan++ {
				pos, length = psa.search(nbuf[scan:])

				for ; scsc < scan+length; scsc++ {
					if scsc+lastoffset < obuflen &&
						obuf[scsc+lastoffset] == nbuf[scsc] {
						oldscore++
					}
				}

				if (length == oldscore && length != 0) || length > oldscore+8 {
					break
				}

				if scan+lastoffset < obuflen && obuf[scan+lastoffset] == nbuf[scan] {
					oldscore--
				}
			}

			if length != oldscore || scan == nbuflen {
				var s, Sf int
				lenf = 0
				for i := int(0); lastscan+i < scan && lastpos+i < obuflen; {
					if obuf[lastpos+i] == nbuf[lastscan+i] {
						s++
					}
					i++
					if s*2-i > Sf*2-lenf {
						Sf = s
						lenf = i
					}
				}

				lenb := 0
				if scan < nbuflen {
					var s, Sb int
					for i := int(1); (scan >= lastscan+i) && (pos >= i); i++ {
						if obuf[pos-i] == nbuf[scan-i] {
							s++
						}
						if s*2-i > Sb*2-lenb {
							Sb = s
							lenb = i
						}
					}
				}

				if lastscan+lenf > scan-lenb {
					overlap := (lastscan + lenf) - (scan - lenb)
					s := int(0)
					Ss := int(0)
					lens := int(0)
					for i := int(0); i < overlap; i++ {
						if nbuf[lastscan+lenf-overlap+i] == obuf[lastpos+lenf-overlap+i] {
							s++
						}
						if nbuf[scan-lenb+i] == obuf[pos-lenb+i] {
							s--
						}
						if s > Ss {
							Ss = s
							lens = i + 1
						}
					}

					lenf += lens - overlap
					lenb -= lens
				}

				c := chunk{
					addOldStart: lastpos,
					addNewStart: lastscan,
					addLength:   lenf,
					copyStart:   lastscan + lenf,
					copyEnd:     scan - lenb,
					offset:      offset,
				}

				if c.addLength > 0 || (c.copyEnd != c.copyStart) {
					// if not a no-op, send
					beforeSend := time.Now()
					chunks <- c
					atomic.AddInt64(&sendingWorkerTime, int64(time.Since(beforeSend)))
				}

				lastscan = scan - lenb
				lastpos = pos - lenb
				lastoffset = pos - scan
			}
		}

		beforeSend := time.Now()
		chunks <- chunk{eoc: true}
		atomic.AddInt64(&sendingWorkerTime, int64(time.Since(beforeSend)))
	}

	blockSize := 256 * 1024
	numBlocks := (nbuflen + blockSize - 1) / blockSize

	if numBlocks < partitions {
		blockSize = nbuflen / partitions
		numBlocks = (nbuflen + blockSize - 1) / blockSize
	}

	numWorkers := partitions * 8

	// fmt.Fprintf(os.Stderr, "Divvying %s in %d block(s) of %s (with %d workers)\n",
	// 	humanize.IBytes(uint64(nbuflen)),
	// 	numBlocks,
	// 	humanize.IBytes(uint64(blockSize)),
	// 	partitions,
	// )

	blockWorkersState := make([]blockWorkerState, numWorkers)

	// initialize all channels
	for i := 0; i < numWorkers; i++ {
		blockWorkersState[i].work = make(chan int, 1)
		blockWorkersState[i].chunks = make(chan chunk, 256)
		blockWorkersState[i].consumed = make(chan bool, 1)
		blockWorkersState[i].consumed <- true
	}

	for i := 0; i < numWorkers; i++ {
		go func(workerState blockWorkerState, workerIndex int) {
			lastWorkTime := time.Now()
			for blockIndex := range workerState.work {
				atomic.AddInt64(&idleWorkerTime, int64(time.Since(lastWorkTime)))

				// fmt.Fprintf(os.Stderr, "\nWorker %d should analyze block %d", workerIndex, blockIndex)
				boundary := blockSize * blockIndex
				realBlockSize := blockSize
				if blockIndex == numBlocks-1 {
					realBlockSize = nbuflen - boundary
				}
				// fmt.Fprintf(os.Stderr, "Analyzing %s block at %d\n", humanize.IBytes(uint64(realBlockSize)), i)

				analyzeBlock(realBlockSize, nbuf[boundary:boundary+realBlockSize], boundary, workerState.chunks)
				// fmt.Fprintf(os.Stderr, "\nWorker %d done analyzing block %d", workerIndex, blockIndex)

				lastWorkTime = time.Now()
			}
		}(blockWorkersState[i], i)
	}

	go func() {
		workerIndex := 0

		for i := 0; i < numBlocks; i++ {
			beforeConsume := time.Now()
			<-blockWorkersState[workerIndex].consumed
			atomic.AddInt64(&waitConsumeTime, int64(time.Since(beforeConsume)))

			beforeEnqueue := time.Now()
			blockWorkersState[workerIndex].work <- i
			atomic.AddInt64(&waitEnqueueTime, int64(time.Since(beforeEnqueue)))

			workerIndex = (workerIndex + 1) % numWorkers
		}

		for workerIndex := 0; workerIndex < numWorkers; workerIndex++ {
			close(blockWorkersState[workerIndex].work)
		}
		// fmt.Fprintf(os.Stderr, "Sent all blockworks\n")
	}()

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan-prepare: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	var prevChunk chunk
	first := true

	consumer.ProgressLabel(fmt.Sprintf("Scanning %s (%d blocks of %s)...", humanize.IBytes(uint64(nbuflen)), numBlocks, humanize.IBytes(uint64(blockSize))))

	workerIndex := 0
	for blockIndex := 0; blockIndex < numBlocks; blockIndex++ {
		// fmt.Fprintf(os.Stderr, "\nWaiting on worker %d for block %d", workerIndex, blockIndex)
		consumer.Progress(float64(blockIndex) / float64(numBlocks))
		state := blockWorkersState[workerIndex]

		for chunk := range state.chunks {
			// fmt.Fprintf(os.Stderr, "\nFor block %d, received chunk %#v", blockIndex, chunk)
			if chunk.eoc {
				break
			}

			if first {
				first = false
			} else {
				bsdc.Seek = int64(chunk.addOldStart - (prevChunk.addOldStart + prevChunk.addLength))

				// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

				err := writeMessage(bsdc)
				if err != nil {
					return err
				}
			}

			ctx.db.Reset()
			ctx.db.Grow(chunk.addLength)

			addNewStart := chunk.addNewStart + chunk.offset

			for i := 0; i < chunk.addLength; i++ {
				ctx.db.WriteByte(nbuf[addNewStart+i] - obuf[chunk.addOldStart+i])
			}

			bsdc.Add = ctx.db.Bytes()
			bsdc.Copy = nbuf[chunk.offset+chunk.copyStart : chunk.offset+chunk.copyEnd]

			if ctx.Stats != nil && ctx.Stats.BiggestAdd < int64(len(bsdc.Add)) {
				ctx.Stats.BiggestAdd = int64(len(bsdc.Add))
			}

			prevChunk = chunk
		}

		state.consumed <- true
		workerIndex = (workerIndex + 1) % numWorkers
	}

	// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

	bsdc.Seek = 0
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentScanning += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		consumer.Debugf("\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	fmt.Fprintf(os.Stderr, "\nStats: waitConsume %s, waitEnqueue %s, idleWorker %s, sendingWorker %s\n",
		time.Duration(waitConsumeTime),
		time.Duration(waitEnqueueTime),
		time.Duration(idleWorkerTime),
		time.Duration(sendingWorkerTime),
	)

	bsdc.Reset()
	bsdc.Eof = true
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>up block size to 1MB<commit_after>package bsdiff

import (
	"fmt"
	"os"
	"runtime"
	"sync/atomic"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/itchio/wharf/state"
)

type chunk struct {
	addOldStart int
	addNewStart int
	addLength   int
	copyStart   int
	copyEnd     int
	offset      int
	eoc         bool
}

type SendChunkFunc func(c chunk)

type blockWorkerState struct {
	consumed chan bool
	work     chan int
	chunks   chan chunk
}

func (ctx *DiffContext) doPartitioned(obuf []byte, obuflen int, nbuf []byte, nbuflen int, memstats *runtime.MemStats, writeMessage WriteMessageFunc, consumer *state.Consumer) error {
	var err error

	partitions := ctx.Partitions
	if partitions >= len(obuf)-1 {
		partitions = 1
	}

	consumer.ProgressLabel(fmt.Sprintf("Sorting %s...", humanize.IBytes(uint64(obuflen))))
	consumer.Progress(0.0)

	startTime := time.Now()

	pmemstats := &runtime.MemStats{}
	runtime.ReadMemStats(pmemstats)
	oldAlloc := pmemstats.TotalAlloc

	if ctx.I == nil {
		ctx.I = make([]int, len(obuf))
		beforeAlloc := time.Now()
		fmt.Fprintf(os.Stderr, "\nAllocated %d-int I in %s\n", len(obuf), time.Since(beforeAlloc))
	} else {
		for len(ctx.I) < len(obuf) {
			lenBefore := len(ctx.I)
			beforeAlloc := time.Now()
			ctx.I = make([]int, len(obuf))
			fmt.Fprintf(os.Stderr, "\nGrown I from %d to %d in %s\n", lenBefore, len(ctx.I), time.Since(beforeAlloc))
		}
	}

	psa := NewPSA(partitions, obuf, ctx.I)

	runtime.ReadMemStats(pmemstats)
	newAlloc := pmemstats.TotalAlloc
	fmt.Fprintf(os.Stderr, "\nAlloc difference after PSA: %s. Size of I: %s\n", humanize.IBytes(uint64(newAlloc-oldAlloc)), humanize.IBytes(uint64(8*len(psa.I))))

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentSorting += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after qsufsort: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	bsdc := &Control{}

	consumer.ProgressLabel(fmt.Sprintf("Preparing to scan %s...", humanize.IBytes(uint64(nbuflen))))
	consumer.Progress(0.0)

	startTime = time.Now()

	var waitConsumeTime int64
	var waitEnqueueTime int64
	var idleWorkerTime int64
	var sendingWorkerTime int64

	analyzeBlock := func(nbuflen int, nbuf []byte, offset int, chunks chan chunk) {
		var lenf int

		// Compute the differences, writing ctrl as we go
		var scan, pos, length int
		var lastscan, lastpos, lastoffset int

		for scan < nbuflen {
			var oldscore int
			scan += length

			for scsc := scan; scan < nbuflen; scan++ {
				pos, length = psa.search(nbuf[scan:])

				for ; scsc < scan+length; scsc++ {
					if scsc+lastoffset < obuflen &&
						obuf[scsc+lastoffset] == nbuf[scsc] {
						oldscore++
					}
				}

				if (length == oldscore && length != 0) || length > oldscore+8 {
					break
				}

				if scan+lastoffset < obuflen && obuf[scan+lastoffset] == nbuf[scan] {
					oldscore--
				}
			}

			if length != oldscore || scan == nbuflen {
				var s, Sf int
				lenf = 0
				for i := int(0); lastscan+i < scan && lastpos+i < obuflen; {
					if obuf[lastpos+i] == nbuf[lastscan+i] {
						s++
					}
					i++
					if s*2-i > Sf*2-lenf {
						Sf = s
						lenf = i
					}
				}

				lenb := 0
				if scan < nbuflen {
					var s, Sb int
					for i := int(1); (scan >= lastscan+i) && (pos >= i); i++ {
						if obuf[pos-i] == nbuf[scan-i] {
							s++
						}
						if s*2-i > Sb*2-lenb {
							Sb = s
							lenb = i
						}
					}
				}

				if lastscan+lenf > scan-lenb {
					overlap := (lastscan + lenf) - (scan - lenb)
					s := int(0)
					Ss := int(0)
					lens := int(0)
					for i := int(0); i < overlap; i++ {
						if nbuf[lastscan+lenf-overlap+i] == obuf[lastpos+lenf-overlap+i] {
							s++
						}
						if nbuf[scan-lenb+i] == obuf[pos-lenb+i] {
							s--
						}
						if s > Ss {
							Ss = s
							lens = i + 1
						}
					}

					lenf += lens - overlap
					lenb -= lens
				}

				c := chunk{
					addOldStart: lastpos,
					addNewStart: lastscan,
					addLength:   lenf,
					copyStart:   lastscan + lenf,
					copyEnd:     scan - lenb,
					offset:      offset,
				}

				if c.addLength > 0 || (c.copyEnd != c.copyStart) {
					// if not a no-op, send
					beforeSend := time.Now()
					chunks <- c
					atomic.AddInt64(&sendingWorkerTime, int64(time.Since(beforeSend)))
				}

				lastscan = scan - lenb
				lastpos = pos - lenb
				lastoffset = pos - scan
			}
		}

		beforeSend := time.Now()
		chunks <- chunk{eoc: true}
		atomic.AddInt64(&sendingWorkerTime, int64(time.Since(beforeSend)))
	}

	blockSize := 1024 * 1024
	numBlocks := (nbuflen + blockSize - 1) / blockSize

	if numBlocks < partitions {
		blockSize = nbuflen / partitions
		numBlocks = (nbuflen + blockSize - 1) / blockSize
	}

	numWorkers := partitions * 8

	// fmt.Fprintf(os.Stderr, "Divvying %s in %d block(s) of %s (with %d workers)\n",
	// 	humanize.IBytes(uint64(nbuflen)),
	// 	numBlocks,
	// 	humanize.IBytes(uint64(blockSize)),
	// 	partitions,
	// )

	blockWorkersState := make([]blockWorkerState, numWorkers)

	// initialize all channels
	for i := 0; i < numWorkers; i++ {
		blockWorkersState[i].work = make(chan int, 1)
		blockWorkersState[i].chunks = make(chan chunk, 256)
		blockWorkersState[i].consumed = make(chan bool, 1)
		blockWorkersState[i].consumed <- true
	}

	for i := 0; i < numWorkers; i++ {
		go func(workerState blockWorkerState, workerIndex int) {
			lastWorkTime := time.Now()
			for blockIndex := range workerState.work {
				atomic.AddInt64(&idleWorkerTime, int64(time.Since(lastWorkTime)))

				// fmt.Fprintf(os.Stderr, "\nWorker %d should analyze block %d", workerIndex, blockIndex)
				boundary := blockSize * blockIndex
				realBlockSize := blockSize
				if blockIndex == numBlocks-1 {
					realBlockSize = nbuflen - boundary
				}
				// fmt.Fprintf(os.Stderr, "Analyzing %s block at %d\n", humanize.IBytes(uint64(realBlockSize)), i)

				analyzeBlock(realBlockSize, nbuf[boundary:boundary+realBlockSize], boundary, workerState.chunks)
				// fmt.Fprintf(os.Stderr, "\nWorker %d done analyzing block %d", workerIndex, blockIndex)

				lastWorkTime = time.Now()
			}
		}(blockWorkersState[i], i)
	}

	go func() {
		workerIndex := 0

		for i := 0; i < numBlocks; i++ {
			beforeConsume := time.Now()
			<-blockWorkersState[workerIndex].consumed
			atomic.AddInt64(&waitConsumeTime, int64(time.Since(beforeConsume)))

			beforeEnqueue := time.Now()
			blockWorkersState[workerIndex].work <- i
			atomic.AddInt64(&waitEnqueueTime, int64(time.Since(beforeEnqueue)))

			workerIndex = (workerIndex + 1) % numWorkers
		}

		for workerIndex := 0; workerIndex < numWorkers; workerIndex++ {
			close(blockWorkersState[workerIndex].work)
		}
		// fmt.Fprintf(os.Stderr, "Sent all blockworks\n")
	}()

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan-prepare: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	var prevChunk chunk
	first := true

	consumer.ProgressLabel(fmt.Sprintf("Scanning %s (%d blocks of %s)...", humanize.IBytes(uint64(nbuflen)), numBlocks, humanize.IBytes(uint64(blockSize))))

	workerIndex := 0
	for blockIndex := 0; blockIndex < numBlocks; blockIndex++ {
		// fmt.Fprintf(os.Stderr, "\nWaiting on worker %d for block %d", workerIndex, blockIndex)
		consumer.Progress(float64(blockIndex) / float64(numBlocks))
		state := blockWorkersState[workerIndex]

		for chunk := range state.chunks {
			// fmt.Fprintf(os.Stderr, "\nFor block %d, received chunk %#v", blockIndex, chunk)
			if chunk.eoc {
				break
			}

			if first {
				first = false
			} else {
				bsdc.Seek = int64(chunk.addOldStart - (prevChunk.addOldStart + prevChunk.addLength))

				// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

				err := writeMessage(bsdc)
				if err != nil {
					return err
				}
			}

			ctx.db.Reset()
			ctx.db.Grow(chunk.addLength)

			addNewStart := chunk.addNewStart + chunk.offset

			for i := 0; i < chunk.addLength; i++ {
				ctx.db.WriteByte(nbuf[addNewStart+i] - obuf[chunk.addOldStart+i])
			}

			bsdc.Add = ctx.db.Bytes()
			bsdc.Copy = nbuf[chunk.offset+chunk.copyStart : chunk.offset+chunk.copyEnd]

			if ctx.Stats != nil && ctx.Stats.BiggestAdd < int64(len(bsdc.Add)) {
				ctx.Stats.BiggestAdd = int64(len(bsdc.Add))
			}

			prevChunk = chunk
		}

		state.consumed <- true
		workerIndex = (workerIndex + 1) % numWorkers
	}

	// fmt.Fprintf(os.Stderr, "%d bytes add, %d bytes copy\n", len(bsdc.Add), len(bsdc.Copy))

	bsdc.Seek = 0
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	if ctx.Stats != nil {
		ctx.Stats.TimeSpentScanning += time.Since(startTime)
	}

	if ctx.MeasureMem {
		runtime.ReadMemStats(memstats)
		consumer.Debugf("\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
		fmt.Fprintf(os.Stderr, "\nAllocated bytes after scan: %s (%s total)", humanize.IBytes(uint64(memstats.Alloc)), humanize.IBytes(uint64(memstats.TotalAlloc)))
	}

	fmt.Fprintf(os.Stderr, "\nStats: waitConsume %s, waitEnqueue %s, idleWorker %s, sendingWorker %s\n",
		time.Duration(waitConsumeTime),
		time.Duration(waitEnqueueTime),
		time.Duration(idleWorkerTime),
		time.Duration(sendingWorkerTime),
	)

	bsdc.Reset()
	bsdc.Eof = true
	err = writeMessage(bsdc)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package protected_objects stores, searches and chains protected objects like keys
// and files.

package protected_objects

import (
	"container/list"
	"errors"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/jlmucb/cloudproxy/go/tpm2"
)

func PrintObject(obj *ObjectMessage) {
	fmt.Printf("Object %s, epoch %d\n", *obj.ObjId.ObjName, *obj.ObjId.ObjEpoch)
	fmt.Printf("\ttype %s, status %s, notbefore: %s, notafter: %s\n", *obj.ObjType,
		*obj.ObjStatus, *obj.NotBefore, *obj.NotAfter)
	fmt.Printf("Object value: %x\n", obj.ObjVal)
}

func PrintProtectedObject(obj *ProtectedObjectMessage) {
	fmt.Printf("Object %s, epoch %d\n", *obj.ProtectedObjId.ObjName, *obj.ProtectedObjId.ObjEpoch)
	fmt.Printf("Object %s, epoch %d\n", *obj.ProtectorObjId.ObjName, *obj.ProtectorObjId.ObjEpoch)
}

func PrintNode(obj *NodeMessage) {
	fmt.Printf("ProtectedObject %s, epoch %d\n", *obj.ProtectedObjId.ObjName,
		 *obj.ProtectedObjId.ObjEpoch)
	fmt.Printf("ProtectorObject %s, epoch %d\n", *obj.ProtectorObjId.ObjName,
		 *obj.ProtectorObjId.ObjEpoch)
}

func CreateObject(name string, epoch int32, obj_type *string, status *string, notBefore *time.Time,
		notAfter *time.Time, v []byte) (*ObjectMessage, error) {
	obj_id := &ObjectIdMessage {
		ObjName: &name,
		ObjEpoch: &epoch,
	}
	str_notBefore := notBefore.String()
	str_notAfter := notAfter.String()
	obj := &ObjectMessage {
		ObjId: obj_id,
		ObjType: obj_type,
		ObjStatus: status,
		NotBefore: &str_notBefore,
		NotAfter: &str_notAfter,
		ObjVal: v,
	}
	return obj, nil
}

func AddObject(l *list.List, obj interface{}) error {
	l.PushFront(obj)
	return nil
}

func DeleteObject(l *list.List, name string, epoch int32) error {
	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ObjectMessage)
		if *o.ObjId.ObjName == name && *o.ObjId.ObjEpoch == epoch {
			l.Remove(e)
			break;
		}
	}
	return nil
}

func DeleteProtectedObject(l *list.List, name string, epoch int32) error {
	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ProtectedObjectMessage)
		if *o.ProtectedObjId.ObjName == name && *o.ProtectedObjId.ObjEpoch == epoch {
			l.Remove(e)
			break;
		}
	}
	return nil
}

func FindProtectedNodes(l *list.List, name string, epoch int32) (*list.List) {
	r := list.New()

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*NodeMessage)
		if epoch != 0 && epoch != *o.ProtectedObjId.ObjEpoch {
			continue
		}
		if name == *o.ProtectedObjId.ObjName {
			r.PushFront(o)
		}
        }
	return r
}

func FindProtectorNodes(l *list.List, name string, epoch int32) (*list.List) {
	r := list.New()

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*NodeMessage)
		if epoch != 0 && epoch != *o.ProtectorObjId.ObjEpoch {
			continue
		}
		if name == *o.ProtectorObjId.ObjName {
			r.PushFront(o)
		}
        }
	return r
}

func FindObject(l *list.List, name string, epoch int32) (*ObjectMessage) {
	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ObjectMessage)
		if epoch != 0 && epoch != *o.ObjId.ObjEpoch {
			continue
		}
		if name == *o.ObjId.ObjName {
			return o
		}
        }
	return nil
}

func GetLatestEpoch(l *list.List, name string, epoch int32) (*ObjectMessage) {
	return nil
}

func GetEarliestEpoch(l *list.List, name string, epoch int32) (*ObjectMessage) {
	return nil
}

func SaveProtectedObjects(l *list.List, file string) error {
	var po_store ProtectedObjectStoreMessage

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ProtectedObjectMessage)
		p := new(ProtectedObjectMessage)
		p.ProtectedObjId.ObjName = o.ProtectedObjId.ObjName
		p.ProtectedObjId.ObjEpoch = o.ProtectedObjId.ObjEpoch
		p.ProtectorObjId.ObjName = o.ProtectorObjId.ObjName
		p.ProtectorObjId.ObjEpoch = o.ProtectorObjId.ObjEpoch
		p.Blob= o.Blob
		po_store.ProtectedObjects = append(po_store.ProtectedObjects, p)
	}
	b, err := proto.Marshal(&po_store)
	if err != nil {
		return err
	}
	ioutil.WriteFile(file, b, 0644)
	return nil
}

func SaveNodes(l *list.List, file string) error {
	var node_store NodeStoreMessage

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*NodeMessage)
		p := new(NodeMessage)
		p.ProtectedObjId.ObjName = o.ProtectedObjId.ObjName
		p.ProtectedObjId.ObjEpoch = o.ProtectedObjId.ObjEpoch
		p.ProtectorObjId.ObjName = o.ProtectorObjId.ObjName
		p.ProtectorObjId.ObjEpoch = o.ProtectorObjId.ObjEpoch
		node_store.NodeObjects = append(node_store.NodeObjects, p)
	}
	b, err := proto.Marshal(&node_store)
	if err != nil {
		return err
	}
	ioutil.WriteFile(file, b, 0644)
	return nil
}

// nil is error return
func SaveObjects(l *list.List, file string) error {
	var o_store ObjectStoreMessage

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(ObjectMessage)
		p := new(ObjectMessage)
		p.ObjId = new(ObjectIdMessage)
		p.ObjId.ObjName = o.ObjId.ObjName
		p.ObjId.ObjEpoch = o.ObjId.ObjEpoch
		p.ObjType = o.ObjType
		p.ObjStatus = o.ObjStatus
		p.NotBefore = o.NotBefore
		p.NotAfter = o.NotAfter
		p.ObjVal = o.ObjVal
		o_store.Objects = append(o_store.Objects, p)
	}
	b, err := proto.Marshal(&o_store)
	if err != nil {
		return nil
	}
	ioutil.WriteFile(file, b, 0644)
	return nil
}

func LoadProtectedObjects(file string) (*list.List) {
	var po_store ProtectedObjectStoreMessage

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	err = proto.Unmarshal(buf, &po_store)
	if err != nil {
		return nil
	}
	l := list.New()
	for _, v := range(po_store.ProtectedObjects) {
		o := new(ProtectedObjectMessage)
		o.ProtectorObjId.ObjName = v.ProtectorObjId.ObjName
		o.ProtectorObjId.ObjEpoch = v.ProtectorObjId.ObjEpoch
		o.ProtectedObjId.ObjName = v.ProtectedObjId.ObjName
		o.ProtectedObjId.ObjEpoch = v.ProtectedObjId.ObjEpoch
		o.Blob = v.Blob
		l.PushFront(o)
	}
	return l
}

func LoadNodes(file string) (*list.List) {
	var node_store NodeStoreMessage

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	err = proto.Unmarshal(buf, &node_store)
	if err != nil {
		return nil
	}
	l := list.New()
	for _, v := range(node_store.NodeObjects) {
		o := new(NodeMessage)
		o.ProtectedObjId.ObjName = v.ProtectedObjId.ObjName
		o.ProtectedObjId.ObjEpoch = v.ProtectedObjId.ObjEpoch
		o.ProtectorObjId.ObjName = v.ProtectorObjId.ObjName
		o.ProtectorObjId.ObjEpoch = v.ProtectorObjId.ObjEpoch
		l.PushFront(o)
	}
	return l
}

func LoadObjects(file string) (*list.List) {
	var o_store ObjectStoreMessage

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	err = proto.Unmarshal(buf, &o_store)
	if err != nil {
		return nil
	}
	l := list.New()
	for _, v := range(o_store.Objects) {
		o := new(ObjectMessage)
		o.ObjId = new(ObjectIdMessage)
		o.ObjId.ObjName = v.ObjId.ObjName
		o.ObjId.ObjEpoch = v.ObjId.ObjEpoch

		o.ObjType = v.ObjType
		o.ObjStatus = v.ObjStatus
		o.NotBefore = v.NotBefore
		o.NotAfter = v.NotAfter
		o.ObjVal = v.ObjVal
		l.PushFront(*o)
	}
	return l
}

func MakeProtectedObject(obj ObjectMessage, protectorName string, protectorEpoch int32,
		protectorKeys []byte) (*ProtectedObjectMessage, error) {
	p := new(ProtectedObjectMessage)
	p.ProtectedObjId = new(ObjectIdMessage)
	p.ProtectorObjId = new(ObjectIdMessage)
	p.ProtectedObjId.ObjName = obj.ObjId.ObjName
	p.ProtectedObjId.ObjEpoch =obj.ObjId.ObjEpoch
	p.ProtectorObjId.ObjName = &protectorName
	p.ProtectorObjId.ObjEpoch = &protectorEpoch
	unencrypted, err := proto.Marshal(&obj)
	if err != nil {
		return nil, errors.New("Can't make Protected Object")
	}
	encrypted, err := tpm2.Protect(protectorKeys, unencrypted)
	if err != nil {
		return nil, errors.New("Can't Protect Object")
	}
	p.Blob = encrypted
	return p, nil
}

func RecoverProtectedObject(obj *ProtectedObjectMessage, protectorKeys []byte) (*ObjectMessage, error) {
	p := new(ObjectMessage)
	unencrypted, err := tpm2.Unprotect(protectorKeys, obj.Blob)
	if err != nil {
		return nil, errors.New("Can't make Unprotect Object")
	}
	err = proto.Unmarshal(unencrypted, p)
	if err != nil {
		return nil, errors.New("Can't Unmarshal Object")
	}
	return p, nil
}

<commit_msg>StoreObjects<commit_after>// Copyright (c) 2014, Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package protected_objects stores, searches and chains protected objects like keys
// and files.

package protected_objects

import (
	"container/list"
	"errors"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/jlmucb/cloudproxy/go/tpm2"
)

func PrintObject(obj *ObjectMessage) {
	fmt.Printf("Object %s, epoch %d\n", *obj.ObjId.ObjName, *obj.ObjId.ObjEpoch)
	fmt.Printf("\ttype %s, status %s, notbefore: %s, notafter: %s\n", *obj.ObjType,
		*obj.ObjStatus, *obj.NotBefore, *obj.NotAfter)
	fmt.Printf("Object value: %x\n", obj.ObjVal)
}

func PrintProtectedObject(obj *ProtectedObjectMessage) {
	fmt.Printf("Object %s, epoch %d\n", *obj.ProtectedObjId.ObjName, *obj.ProtectedObjId.ObjEpoch)
	fmt.Printf("Object %s, epoch %d\n", *obj.ProtectorObjId.ObjName, *obj.ProtectorObjId.ObjEpoch)
}

func PrintNode(obj *NodeMessage) {
	fmt.Printf("ProtectedObject %s, epoch %d\n", *obj.ProtectedObjId.ObjName,
		 *obj.ProtectedObjId.ObjEpoch)
	fmt.Printf("ProtectorObject %s, epoch %d\n", *obj.ProtectorObjId.ObjName,
		 *obj.ProtectorObjId.ObjEpoch)
}

func CreateObject(name string, epoch int32, obj_type *string, status *string, notBefore *time.Time,
		notAfter *time.Time, v []byte) (*ObjectMessage, error) {
	obj_id := &ObjectIdMessage {
		ObjName: &name,
		ObjEpoch: &epoch,
	}
	str_notBefore := notBefore.String()
	str_notAfter := notAfter.String()
	obj := &ObjectMessage {
		ObjId: obj_id,
		ObjType: obj_type,
		ObjStatus: status,
		NotBefore: &str_notBefore,
		NotAfter: &str_notAfter,
		ObjVal: v,
	}
	return obj, nil
}

func AddObject(l *list.List, obj interface{}) error {
	l.PushFront(obj)
	return nil
}

func DeleteObject(l *list.List, name string, epoch int32) error {
	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ObjectMessage)
		if *o.ObjId.ObjName == name && *o.ObjId.ObjEpoch == epoch {
			l.Remove(e)
			break;
		}
	}
	return nil
}

func DeleteProtectedObject(l *list.List, name string, epoch int32) error {
	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ProtectedObjectMessage)
		if *o.ProtectedObjId.ObjName == name && *o.ProtectedObjId.ObjEpoch == epoch {
			l.Remove(e)
			break;
		}
	}
	return nil
}

func FindProtectedNodes(l *list.List, name string, epoch int32) (*list.List) {
	r := list.New()

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*NodeMessage)
		if epoch != 0 && epoch != *o.ProtectedObjId.ObjEpoch {
			continue
		}
		if name == *o.ProtectedObjId.ObjName {
			r.PushFront(o)
		}
        }
	return r
}

func FindProtectorNodes(l *list.List, name string, epoch int32) (*list.List) {
	r := list.New()

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*NodeMessage)
		if epoch != 0 && epoch != *o.ProtectorObjId.ObjEpoch {
			continue
		}
		if name == *o.ProtectorObjId.ObjName {
			r.PushFront(o)
		}
        }
	return r
}

func FindObject(l *list.List, name string, epoch int32) (*ObjectMessage) {
	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(*ObjectMessage)
		if epoch != 0 && epoch != *o.ObjId.ObjEpoch {
			continue
		}
		if name == *o.ObjId.ObjName {
			return o
		}
        }
	return nil
}

func GetLatestEpoch(l *list.List, name string, epoch int32) (*ObjectMessage) {
	return nil
}

func GetEarliestEpoch(l *list.List, name string, epoch int32) (*ObjectMessage) {
	return nil
}

func SaveProtectedObjects(l *list.List, file string) error {
	var po_store ProtectedObjectStoreMessage

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(ProtectedObjectMessage)
		p := new(ProtectedObjectMessage)
		p.ProtectedObjId.ObjName = o.ProtectedObjId.ObjName
		p.ProtectedObjId.ObjEpoch = o.ProtectedObjId.ObjEpoch
		p.ProtectorObjId.ObjName = o.ProtectorObjId.ObjName
		p.ProtectorObjId.ObjEpoch = o.ProtectorObjId.ObjEpoch
		p.Blob= o.Blob
		po_store.ProtectedObjects = append(po_store.ProtectedObjects, p)
	}
	b, err := proto.Marshal(&po_store)
	if err != nil {
		return err
	}
	ioutil.WriteFile(file, b, 0644)
	return nil
}

func SaveNodes(l *list.List, file string) error {
	var node_store NodeStoreMessage

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(NodeMessage)
		p := new(NodeMessage)
		p.ProtectedObjId.ObjName = o.ProtectedObjId.ObjName
		p.ProtectedObjId.ObjEpoch = o.ProtectedObjId.ObjEpoch
		p.ProtectorObjId.ObjName = o.ProtectorObjId.ObjName
		p.ProtectorObjId.ObjEpoch = o.ProtectorObjId.ObjEpoch
		node_store.NodeObjects = append(node_store.NodeObjects, p)
	}
	b, err := proto.Marshal(&node_store)
	if err != nil {
		return err
	}
	ioutil.WriteFile(file, b, 0644)
	return nil
}

// nil is error return
func SaveObjects(l *list.List, file string) error {
	var o_store ObjectStoreMessage

	for e := l.Front(); e != nil; e = e.Next() {
		o := e.Value.(ObjectMessage)
		p := new(ObjectMessage)
		p.ObjId = new(ObjectIdMessage)
		p.ObjId.ObjName = o.ObjId.ObjName
		p.ObjId.ObjEpoch = o.ObjId.ObjEpoch
		p.ObjType = o.ObjType
		p.ObjStatus = o.ObjStatus
		p.NotBefore = o.NotBefore
		p.NotAfter = o.NotAfter
		p.ObjVal = o.ObjVal
		o_store.Objects = append(o_store.Objects, p)
	}
	b, err := proto.Marshal(&o_store)
	if err != nil {
		return nil
	}
	ioutil.WriteFile(file, b, 0644)
	return nil
}

func LoadProtectedObjects(file string) (*list.List) {
	var po_store ProtectedObjectStoreMessage

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	err = proto.Unmarshal(buf, &po_store)
	if err != nil {
		return nil
	}
	l := list.New()
	for _, v := range(po_store.ProtectedObjects) {
		o := new(ProtectedObjectMessage)
		o.ProtectorObjId.ObjName = v.ProtectorObjId.ObjName
		o.ProtectorObjId.ObjEpoch = v.ProtectorObjId.ObjEpoch
		o.ProtectedObjId.ObjName = v.ProtectedObjId.ObjName
		o.ProtectedObjId.ObjEpoch = v.ProtectedObjId.ObjEpoch
		o.Blob = v.Blob
		l.PushFront(*o)
	}
	return l
}

func LoadNodes(file string) (*list.List) {
	var node_store NodeStoreMessage

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	err = proto.Unmarshal(buf, &node_store)
	if err != nil {
		return nil
	}
	l := list.New()
	for _, v := range(node_store.NodeObjects) {
		o := new(NodeMessage)
		o.ProtectedObjId.ObjName = v.ProtectedObjId.ObjName
		o.ProtectedObjId.ObjEpoch = v.ProtectedObjId.ObjEpoch
		o.ProtectorObjId.ObjName = v.ProtectorObjId.ObjName
		o.ProtectorObjId.ObjEpoch = v.ProtectorObjId.ObjEpoch
		l.PushFront(*o)
	}
	return l
}

func LoadObjects(file string) (*list.List) {
	var o_store ObjectStoreMessage

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	err = proto.Unmarshal(buf, &o_store)
	if err != nil {
		return nil
	}
	l := list.New()
	for _, v := range(o_store.Objects) {
		o := new(ObjectMessage)
		PrintObject(v)
		o.ObjId = new(ObjectIdMessage)
		o.ObjId.ObjName = v.ObjId.ObjName
		o.ObjId.ObjEpoch = v.ObjId.ObjEpoch

		o.ObjType = v.ObjType
		o.ObjStatus = v.ObjStatus
		o.NotBefore = v.NotBefore
		o.NotAfter = v.NotAfter
		o.ObjVal = v.ObjVal
		l.PushFront(*o)
	}
	return l
}

func MakeProtectedObject(obj ObjectMessage, protectorName string, protectorEpoch int32,
		protectorKeys []byte) (*ProtectedObjectMessage, error) {
	p := new(ProtectedObjectMessage)
	p.ProtectedObjId = new(ObjectIdMessage)
	p.ProtectorObjId = new(ObjectIdMessage)
	p.ProtectedObjId.ObjName = obj.ObjId.ObjName
	p.ProtectedObjId.ObjEpoch =obj.ObjId.ObjEpoch
	p.ProtectorObjId.ObjName = &protectorName
	p.ProtectorObjId.ObjEpoch = &protectorEpoch
	unencrypted, err := proto.Marshal(&obj)
	if err != nil {
		return nil, errors.New("Can't make Protected Object")
	}
	encrypted, err := tpm2.Protect(protectorKeys, unencrypted)
	if err != nil {
		return nil, errors.New("Can't Protect Object")
	}
	p.Blob = encrypted
	return p, nil
}

func RecoverProtectedObject(obj *ProtectedObjectMessage, protectorKeys []byte) (*ObjectMessage, error) {
	p := new(ObjectMessage)
	unencrypted, err := tpm2.Unprotect(protectorKeys, obj.Blob)
	if err != nil {
		return nil, errors.New("Can't make Unprotect Object")
	}
	err = proto.Unmarshal(unencrypted, p)
	if err != nil {
		return nil, errors.New("Can't Unmarshal Object")
	}
	return p, nil
}

<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	. "koding/db/models"
	helper "koding/db/mongodb/modelhelper"
)

var POST_TYPES = [5]string{
	"JBlogPost",
	"JDiscussion",
	"JCodeSnip",
	"JTutorial",
	"JStatusUpdate",
}

const LIMIT = 100

type JPost struct {
	Title        string `bson:"title,omitempty"`
	OpinionCount int    `bson:"opinionCount,omitempty"`
}

var ErrAlreadyMigrated = errors.New("already migrated")

func main() {
	initPublisher()
	defer shutdown()

	for _, postType := range POST_TYPES {
		log.Notice("Starting \"%s\" migration", postType)
		m := &Migrator{
			PostType: postType,
		}
		if err := initialize(m); err != nil {
			log.Error("An error occured during migration: %s", err.Error())
		}
		GetMigrationCompletedReport(m)
	}
}

func initialize(m *Migrator) error {
	count, err := helper.CountPosts(helper.Selector{}, helper.Options{}, m.PostType)
	if err != nil {
		return err
	}
	log.Info("Found %v posts", count)

	return migrate(m)
}

func migrate(m *Migrator) error {
	o := helper.Options{
		Sort:  "meta.createdAt", //start from the oldest
		Limit: LIMIT,
		Skip:  m.Index,
	}
	posts, err := helper.GetSomePosts(helper.Selector{}, o, m.PostType)
	if err != nil {
		return err
	}
	// no more post to migrate
	if len(posts) == 0 {
		return nil
	}

	for _, post := range posts {
		m.Id = post.Id.Hex()
		post.Id = helper.NewObjectId()
		m.NewId = post.Id.Hex()
		m.Index++
		if err := verifyOrigin(&post); err != nil {
			ReportError(m, err)
			continue
		}

		if err := insertNewStatusUpdate(&post, m); err != nil {
			ReportError(m, err)
			continue
		}

		if err := migrateTags(&post, m); err != nil {
			ReportError(m, err)
			continue
		}

		if err := migrateOrigin(&post); err != nil {
			ReportError(m, err)
			continue
		}

		commenters, err := migrateComments(&post, m)
		if err != nil {
			ReportError(m, err)
			continue
		}

		if err := migrateOpinions(&post, m, commenters); err != nil {
			ReportError(m, err)
			continue
		}

		ReportSuccess(m)
	}
	return migrate(m)
}

func insertNewStatusUpdate(p *Post, m *Migrator) error {
	if p.MigrationStatus == "Completed" {
		return ErrAlreadyMigrated
	}
	exists, err := helper.CheckGroupExistence(p.Group)
	if err != nil {
		return err
	}
	if !exists {
		return fmt.Errorf("Group \"%s\" not found", p.Group)
	}
	// some status update have negative timestamp
	if p.Meta.CreatedAt.Unix() < 0 {
		return fmt.Errorf("got nil timestamp for")
	}
	p.Meta.Likes = 0 // CtF: because we are not migrating activities it is reset

	if err := migrateCodesnip(p, m); err != nil {
		return err
	}

	su := p.ConvertToStatusUpdate()
	return helper.AddStatusUpdate(su)
}

func migrateCodesnip(p *Post, m *Migrator) error {
	if m.PostType != "JCodeSnip" {
		return nil
	}

	if len(p.Attachments) > 0 {
		body, ok := p.Attachments[0]["content"]
		if !ok {
			return fmt.Errorf("Codesnip content not found")
		}
		p.Body = fmt.Sprintf("`%s`", body)
	}
	return nil
}

func migrateTags(p *Post, m *Migrator) error {
	s := helper.Selector{
		"sourceId":   helper.GetObjectId(m.Id),
		"as":         "tag",
		"targetName": "JTag",
	}
	rels, err := helper.GetAllRelationships(s)
	if err != nil {
		return err
	}
	log.Info("%v tags found", len(rels))

	for _, r := range rels {
		tagId := r.TargetId.Hex()
		// first check tag existence
		exists, err := helper.CheckTagExistence(tagId)
		if err != nil {
			return err
		}
		if !exists {
			continue
		}

		// update tag relationships
		r.Id = helper.NewObjectId()
		r.SourceId = p.Id
		r.SourceName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			return err
		}
		sr := swapRelationship(&r, "post")
		if err := helper.AddRelationship(sr); err != nil {
			return err
		}
		//append tag to status update body
		p.Body += fmt.Sprintf(" |#:JTag:%s|", tagId)
	}
	// if post is tagged
	if len(rels) > 0 {
		su := p.ConvertToStatusUpdate()
		return helper.UpdateStatusUpdate(su)
	}

	return nil
}

func migrateComments(p *Post, m *Migrator) (map[string]bool, error) {
	accounts := make(map[string]bool)
	// get all comments
	s := helper.Selector{
		"sourceId":   helper.GetObjectId(m.Id),
		"as":         "reply",
		"targetName": "JComment",
	}
	rels, err := helper.GetAllRelationships(s)
	if err != nil {
		return accounts, err
	}
	log.Info("%v comments found", len(rels))
	// posts does not have any comments
	if len(rels) == 0 {
		return accounts, nil
	}

	count := 0

	for _, r := range rels {
		comment, err := helper.GetCommentById(r.TargetId.Hex())
		if err != nil {
			if err == helper.ErrNotFound {
				log.Info("Comment not found - Id: %s", r.TargetId.Hex())
				continue
			}
			return accounts, err
		}

		// check origin existence
		originId := comment.OriginId.Hex()
		originExists, err := helper.CheckAccountExistence(originId)
		if err != nil {
			return accounts, err
		}
		if !originExists {
			continue
		}

		// add relationship: JNewStatusUpdate -> reply -> JComment
		r.Id = helper.NewObjectId()
		r.SourceId = p.Id
		r.SourceName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			return accounts, err
		}

		// get unique commenters for each post
		if _, exist := accounts[originId]; !exist {
			accounts[originId] = true
			migrateCommentOrigin(comment, p)
		}
		count++
	}

	p.RepliesCount = count

	return accounts, nil
}

// migrateOpinions migrates opinions to comments
// JDiscussion opinion JOpinion
// JAccount creator JOpinion
func migrateOpinions(p *Post, m *Migrator, commenters map[string]bool) error {
	if m.PostType != "JDiscussion" && m.PostType != "JTutorial" {
		return nil
	}
	s := helper.Selector{
		"sourceId":   helper.GetObjectId(m.Id),
		"as":         "opinion",
		"targetName": "JOpinion",
	}
	rels, err := helper.GetAllRelationships(s)
	if err != nil {
		return err
	}
	log.Info("%v opinions found", len(rels))
	// post does not have any opinion
	if len(rels) == 0 {
		return nil
	}
	count := 0
	for _, r := range rels {
		opinion, err := helper.GetOpinionById(r.TargetId.Hex())
		if err != nil {
			if err == helper.ErrNotFound {
				log.Info("Opinion not found - Id: %s", r.TargetId.Hex())
				continue
			}
			return err
		}
		// check origin existence
		originId := opinion.OriginId.Hex()
		originExists, err := helper.CheckAccountExistence(originId)
		if err != nil {
			return err
		}
		if !originExists {
			continue
		}
		comment, err := convertOpinionToComment(opinion, p)
		if err != nil {
			return err
		}

		if _, exist := commenters[originId]; !exist {
			commenters[originId] = true
			if err := migrateCommentOrigin(comment, opinion); err != nil {
				return err
			}
		}

		count++
	}
	p.RepliesCount += count
	return nil
}

// convertOpinionToComment first converts opinion to comment and persists it.
// Adds Relationship: JNewStatusUpdate -> reply    -> JComment
//                    JAccount         -> creator -> JComment
func convertOpinionToComment(opinion *Post, post *Post) (*Comment, error) {
	c := &Comment{
		Id:         helper.NewObjectId(),
		Body:       opinion.Body,
		OriginType: opinion.OriginType,
		OriginId:   opinion.OriginId,
		Meta:       opinion.Meta,
	}
	c.Meta.Likes = 0 // TODO not sure about it
	if err := helper.AddComment(c); err != nil {
		return c, err
	}

	// Add relationship: JNewStatusUpdate -> reply -> JComment
	r := &Relationship{
		Id:         helper.NewObjectId(),
		SourceId:   post.Id,
		SourceName: "JNewStatusUpdate",
		TargetId:   c.Id,
		TargetName: "JComment",
		As:         "reply",
		TimeStamp:  opinion.Meta.CreatedAt,
	}
	if err := helper.AddRelationship(r); err != nil {
		return c, err
	}
	return c, addCommentCreator(c, r)
}

// addCommentCreator inserts a new relationship as JAccount -> creator -> JComment
func addCommentCreator(c *Comment, r *Relationship) error {
	r.Id = helper.NewObjectId()
	r.SourceId = c.OriginId
	r.SourceName = "JAccount"
	r.As = "creator"
	r.TimeStamp = c.Meta.CreatedAt
	return helper.AddRelationship(r)
}

// migrateCommentOrigins inserts commenter and follower relationships
// JNewStatusUpdate -> commenter -> JAccount
// JNewStatusUpdate -> follower -> JAccount
func migrateCommentOrigin(c *Comment, p *Post) error {
	r := &Relationship{
		Id:         helper.NewObjectId(),
		SourceId:   p.Id,
		SourceName: "JNewStatusUpdate",
		TargetId:   c.OriginId,
		TargetName: "JAccount",
		As:         "commenter",
		TimeStamp:  c.Meta.CreatedAt,
	}
	if err := helper.AddRelationship(r); err != nil {
		return err
	}
	r.Id = helper.NewObjectId()
	r.As = "follower"
	return helper.AddRelationship(r)
}

func migrateOrigin(p *Post) error {
	r := &Relationship{
		Id:         helper.NewObjectId(),
		SourceId:   p.OriginId,
		SourceName: "JAccount",
		TargetId:   p.Id,
		TargetName: "JNewStatusUpdate",
		TimeStamp:  p.Meta.CreatedAt,
		As:         "creator",
	}
	if err := helper.AddRelationship(r); err != nil {
		return err
	}

	r = swapRelationship(r, "author")
	return helper.AddRelationship(r)
}

func verifyOrigin(p *Post) error {
	originId := p.OriginId.Hex()
	if originId == "" {
		return fmt.Errorf("Empty origin id")
	}

	result, err := helper.CheckAccountExistence(originId)
	if err != nil {
		return err
	}
	if !result {
		return fmt.Errorf("Origin not found - %v", originId)
	}

	return nil
}

// swapTagRelation swaps source and target data of relationships. It is used
// for converting bidirectional relationships.
func swapRelationship(r *Relationship, as string) *Relationship {
	return &Relationship{
		Id:         helper.NewObjectId(),
		As:         as,
		SourceId:   r.TargetId,
		SourceName: r.TargetName,
		TargetId:   r.SourceId,
		TargetName: r.SourceName,
		TimeStamp:  r.TimeStamp,
		Data:       r.Data,
	}
}
<commit_msg>Migration: migration can be started from the last interruption point<commit_after>package main

import (
	"errors"
	"fmt"
	. "koding/db/models"
	helper "koding/db/mongodb/modelhelper"
)

var POST_TYPES = [5]string{
	"JBlogPost",
	"JDiscussion",
	"JCodeSnip",
	"JTutorial",
	"JStatusUpdate",
}

const LIMIT = 100

type JPost struct {
	Title        string `bson:"title,omitempty"`
	OpinionCount int    `bson:"opinionCount,omitempty"`
}

var ErrAlreadyMigrated = errors.New("already migrated")

func main() {
	initPublisher()
	defer shutdown()

	for _, postType := range POST_TYPES {
		log.Notice("Starting \"%s\" migration", postType)
		m := &Migrator{
			PostType: postType,
		}
		if err := initialize(m); err != nil {
			log.Error("An error occured during migration: %s", err.Error())
		}
		GetMigrationCompletedReport(m)
	}
}

func initialize(m *Migrator) error {
	count, err := helper.CountPosts(helper.Selector{}, helper.Options{}, m.PostType)
	if err != nil {
		return err
	}
	log.Info("Found %v posts", count)

	return migrate(m)
}

func migrate(m *Migrator) error {
	o := helper.Options{
		Sort:  "meta.createdAt", //start from the oldest
		Limit: LIMIT,
		Skip:  m.Index,
	}
	posts, err := helper.GetSomePosts(helper.Selector{}, o, m.PostType)
	if err != nil {
		return err
	}
	// no more post to migrate
	if len(posts) == 0 {
		return nil
	}

	for _, post := range posts {
		m.Id = post.Id.Hex()
		oldPost := post
		post.Id = helper.NewObjectId()
		m.NewId = post.Id.Hex()
		m.Index++
		if err := verifyOrigin(&post); err != nil {
			ReportError(m, err)
			continue
		}

		if err := insertNewStatusUpdate(&post, m); err != nil {
			ReportError(m, err)
			continue
		}

		if err := updatePostStatus(&oldPost, m, "Started"); err != nil {
			ReportError(m, err)
			continue
		}

		if err := migrateTags(&post, m); err != nil {
			ReportError(m, err)
			continue
		}

		if err := migrateOrigin(&post, m); err != nil {
			ReportError(m, err)
			continue
		}

		commenters, err := migrateComments(&post, m)
		if err != nil {
			ReportError(m, err)
			continue
		}
		if err := migrateOpinions(&post, m, commenters); err != nil {
			ReportError(m, err)
			continue
		}

		if err := updatePostStatus(&oldPost, m, "Completed"); err != nil {
			ReportError(m, err)
			continue
		}

		// updates repliesCount
		su := post.ConvertToStatusUpdate()
		su.MigrationStatus = m.PostType
		if err := helper.UpdateStatusUpdate(su); err != nil {
			ReportError(m, err)
			continue
		}

		ReportSuccess(m)
	}
	return migrate(m)
}

func updatePostStatus(p *Post, m *Migrator, status string) error {
	p.MigrationStatus = status
	return helper.UpdatePost(p, m.PostType)
}

func updateRelationshipStatus(r Relationship, status string) {
	r.MigrationStatus = status
	helper.UpdateRelationship(&r)
}

func insertNewStatusUpdate(p *Post, m *Migrator) error {
	if p.MigrationStatus == "Completed" {
		return ErrAlreadyMigrated
	}
	// it seems post is already migrated with some incomplete relationships
	if p.MigrationStatus == "Started" {
		ep, err := helper.GetStatusUpdate(helper.Selector{"slug": p.Slug})
		p.Id = ep.Id
		return err
	}
	exists, err := helper.CheckGroupExistence(p.Group)
	if err != nil {
		return err
	}
	if !exists {
		return fmt.Errorf("Group \"%s\" not found", p.Group)
	}
	// some status update have negative timestamp
	if p.Meta.CreatedAt.Unix() < 0 {
		return fmt.Errorf("got nil timestamp for")
	}
	p.Meta.Likes = 0 // CtF: because we are not migrating activities it is reset

	if err := migrateCodesnip(p, m); err != nil {
		return err
	}

	su := p.ConvertToStatusUpdate()
	return helper.AddStatusUpdate(su)
}

func migrateCodesnip(p *Post, m *Migrator) error {
	if m.PostType != "JCodeSnip" {
		return nil
	}

	if len(p.Attachments) > 0 {
		codesnip, ok := p.Attachments[0]["content"]
		if !ok {
			return fmt.Errorf("Codesnip content not found")
		}
		// concatenate post body with codesnip
		p.Body = fmt.Sprintf("%s \n\n ```%s``` \n\n", p.Body, codesnip)
		p.Attachments = make([]map[string]interface{}, 0)
	}
	return nil
}

func migrateTags(p *Post, m *Migrator) error {
	s := helper.Selector{
		"sourceId":   helper.GetObjectId(m.Id),
		"as":         "tag",
		"targetName": "JTag",
	}
	rels, err := helper.GetAllRelationships(s)
	if err != nil {
		return err
	}
	log.Info("%v tags found", len(rels))

	for _, r := range rels {
		// relationship already migrated
		if r.MigrationStatus == "Completed" {
			continue
		}
		tagId := r.TargetId.Hex()
		// first check tag existence
		exists, err := helper.CheckTagExistence(tagId)
		if err != nil {
			updateRelationshipStatus(r, "Error")
			return err
		}
		if !exists {
			updateRelationshipStatus(r, "Error")
			continue
		}
		or := r // copy old relationship
		// update tag relationships
		r.Id = helper.NewObjectId()
		r.SourceId = p.Id
		r.SourceName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
		sr := swapRelationship(&r, "post") // CtF: leaking relationship
		if err := helper.AddRelationship(sr); err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
		//append tag to status update body
		p.Body += fmt.Sprintf(" |#:JTag:%s|", tagId)
		updateRelationshipStatus(or, "Completed")
	}

	return nil
}

func migrateComments(p *Post, m *Migrator) (map[string]bool, error) {
	accounts := make(map[string]bool)
	// get all comments
	s := helper.Selector{
		"sourceId":   helper.GetObjectId(m.Id),
		"as":         "reply",
		"targetName": "JComment",
	}
	rels, err := helper.GetAllRelationships(s)
	if err != nil {
		return accounts, err
	}
	log.Info("%v comments found", len(rels))
	// posts does not have any comments
	if len(rels) == 0 {
		return accounts, nil
	}

	count := 0

	for _, r := range rels {
		or := r
		comment, err := helper.GetCommentById(r.TargetId.Hex())
		if err != nil {
			updateRelationshipStatus(or, "Error")
			if err == helper.ErrNotFound {
				log.Info("Comment not found - Id: %s", r.TargetId.Hex())
				continue
			}
			return accounts, err
		}

		// check origin existence
		originId := comment.OriginId.Hex()
		originExists, err := helper.CheckAccountExistence(originId)
		if err != nil {
			updateRelationshipStatus(or, "Error")
			return accounts, err
		}
		if !originExists {
			updateRelationshipStatus(or, "Error")
			continue
		}
		if r.MigrationStatus != "Completed" {
			// add relationship: JNewStatusUpdate -> reply -> JComment
			r.Id = helper.NewObjectId()
			r.SourceId = p.Id
			r.SourceName = "JNewStatusUpdate"
			if err := helper.AddRelationship(&r); err != nil {
				updateRelationshipStatus(or, "Error")
				return accounts, err
			}

			// get unique commenters for each post
			if _, exist := accounts[originId]; !exist {
				accounts[originId] = true
				if err := migrateCommentOrigin(comment, p, m); err != nil {
					return accounts, err
				}
			}
			updateRelationshipStatus(or, "Completed")
		} else {
			accounts[originId] = true
		}

		count++

	}

	p.RepliesCount = count

	return accounts, nil
}

// migrateOpinions migrates opinions to comments
// JDiscussion opinion JOpinion
// JAccount creator JOpinion
func migrateOpinions(p *Post, m *Migrator, commenters map[string]bool) error {
	if m.PostType != "JDiscussion" && m.PostType != "JTutorial" {
		return nil
	}
	s := helper.Selector{
		"sourceId":   helper.GetObjectId(m.Id),
		"as":         "opinion",
		"targetName": "JOpinion",
	}
	rels, err := helper.GetAllRelationships(s)
	if err != nil {
		return err
	}
	log.Info("%v opinions found", len(rels))
	// post does not have any opinion
	if len(rels) == 0 {
		return nil
	}
	count := 0
	for _, r := range rels {
		or := r
		opinion, err := helper.GetOpinionById(r.TargetId.Hex())
		if err != nil {
			updateRelationshipStatus(or, "Error")
			if err == helper.ErrNotFound {
				log.Info("Opinion not found - Id: %s", r.TargetId.Hex())
				continue
			}
			return err
		}
		// check origin existence
		originId := opinion.OriginId.Hex()
		originExists, err := helper.CheckAccountExistence(originId)
		if err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
		if !originExists {
			continue
		}
		if r.MigrationStatus != "Completed" {
			comment, err := convertOpinionToComment(opinion, p)
			if err != nil {
				updateRelationshipStatus(or, "Error")
				return err
			}
			if _, exist := commenters[originId]; !exist {
				commenters[originId] = true
				if err := migrateCommentOrigin(comment, p, m); err != nil {
					return err
				}
			}
			updateRelationshipStatus(or, "Completed")
		} else {
			commenters[originId] = true
		}

		count++
	}
	p.RepliesCount += count
	return nil
}

// convertOpinionToComment first converts opinion to comment and persists it.
// Adds Relationship: JNewStatusUpdate -> reply    -> JComment
//                    JAccount         -> creator -> JComment
func convertOpinionToComment(opinion *Post, post *Post) (*Comment, error) {
	c := &Comment{
		Id:         helper.NewObjectId(),
		Body:       opinion.Body,
		OriginType: opinion.OriginType,
		OriginId:   opinion.OriginId,
		Meta:       opinion.Meta,
	}
	c.Meta.Likes = 0 // TODO not sure about it
	if err := helper.AddComment(c); err != nil {
		return c, err
	}

	// Add relationship: JNewStatusUpdate -> reply -> JComment
	r := &Relationship{
		Id:         helper.NewObjectId(),
		SourceId:   post.Id,
		SourceName: "JNewStatusUpdate",
		TargetId:   c.Id,
		TargetName: "JComment",
		As:         "reply",
		TimeStamp:  opinion.Meta.CreatedAt,
	}
	if err := helper.AddRelationship(r); err != nil {
		return c, err
	}
	return c, addCommentCreator(c, r)
}

// addCommentCreator inserts a new relationship as JAccount -> creator -> JComment
func addCommentCreator(c *Comment, r *Relationship) error {
	r.Id = helper.NewObjectId()
	r.SourceId = c.OriginId
	r.SourceName = "JAccount"
	r.As = "creator"
	r.TimeStamp = c.Meta.CreatedAt
	return helper.AddRelationship(r)
}

// migrateCommentOrigins inserts commenter and follower relationships
// JNewStatusUpdate -> commenter -> JAccount
// JNewStatusUpdate -> follower -> JAccount
func migrateCommentOrigin(c *Comment, p *Post, m *Migrator) error {
	s := helper.Selector{
		"sourceId": helper.GetObjectId(m.Id),
		"targetId": c.OriginId,
		"as":       "commenter",
	}
	r, err := helper.GetRelationship(s)
	if err != nil {
		return fmt.Errorf("commenter not found")
	}
	if r.MigrationStatus != "Completed" {
		or := r
		r.Id = helper.NewObjectId()
		r.SourceId = p.Id
		r.SourceName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
	}
	s["as"] = "follower"
	r, err = helper.GetRelationship(s)
	if err != nil {
		return fmt.Errorf("follower not found")
	}

	if r.MigrationStatus != "Completed" {
		or := r
		r.Id = helper.NewObjectId()
		r.SourceId = p.Id
		r.SourceName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
	}

	return nil
}

func migrateOrigin(p *Post, m *Migrator) error {
	s := helper.Selector{
		"as":       "creator",
		"targetId": helper.GetObjectId(m.Id),
		"sourceId": p.OriginId,
	}
	r, err := helper.GetRelationship(s)
	if err != nil {
		return fmt.Errorf("creator not found")
	}

	if r.MigrationStatus != "Completed" {
		or := r
		r.Id = helper.NewObjectId()
		r.TargetId = p.Id
		r.TargetName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
		updateRelationshipStatus(or, "Completed")
	}
	s = helper.Selector{
		"as":       "author",
		"sourceId": helper.GetObjectId(m.Id),
		"targetId": p.OriginId,
	}
	r, err = helper.GetRelationship(s)
	if err != nil {
		return fmt.Errorf("author not found")
	}

	if r.MigrationStatus != "Completed" {
		or := r
		r.Id = helper.NewObjectId()
		r.SourceId = p.Id
		r.SourceName = "JNewStatusUpdate"
		if err := helper.AddRelationship(&r); err != nil {
			updateRelationshipStatus(or, "Error")
			return err
		}
		updateRelationshipStatus(or, "Completed")
	}

	return nil
}

func verifyOrigin(p *Post) error {
	originId := p.OriginId.Hex()
	if originId == "" {
		return fmt.Errorf("Empty origin id")
	}

	result, err := helper.CheckAccountExistence(originId)
	if err != nil {
		return err
	}
	if !result {
		return fmt.Errorf("Origin not found - %v", originId)
	}

	return nil
}

// swapTagRelation swaps source and target data of relationships. It is used
// for converting bidirectional relationships.
func swapRelationship(r *Relationship, as string) *Relationship {
	return &Relationship{
		Id:         helper.NewObjectId(),
		As:         as,
		SourceId:   r.TargetId,
		SourceName: r.TargetName,
		TargetId:   r.SourceId,
		TargetName: r.SourceName,
		TimeStamp:  r.TimeStamp,
		Data:       r.Data,
	}
}
<|endoftext|>