<commit_before>package scheduler

import (
	"log"

	"github.com/chrislusf/glow/flow"
)

type TaskGroupStatus int

const (
	New TaskGroupStatus = iota
	Waiting
	Scheduled
	Completed
	Failed
)

type TaskGroup struct {
	Id              int
	Tasks           []*flow.Task
	Parents         []*TaskGroup
	ParentStepGroup *StepGroup
	Status          TaskGroupStatus
}

type StepGroup struct {
	Steps      []*flow.Step
	Parents    []*StepGroup
	TaskGroups []*TaskGroup
}

func GroupTasks(fc *flow.FlowContext) []*TaskGroup {
	stepGroups := translateToStepGroups(fc)
	return translateToTaskGroups(stepGroups)
}

func isMergeableDataset(ds *flow.Dataset, taskCount int) bool {
	if taskCount != len(ds.Shards) {
		return false
	}
	if taskCount != len(ds.Step.Tasks) {
		return false
	}
	if len(ds.ReadingSteps) > 1 {
		return false
	}
	for _, shard := range ds.Shards {
		if len(shard.ReadingTasks) > 1 {
			return false
		}
	}
	return true
}

// find mergeable parent step or itself if parent is not mergeable
func findAncestorStepId(step *flow.Step) (int, bool) {
	current := step
	taskCount := len(current.Tasks)

	// println("find step", step.Name)

	for taskCount == len(current.Tasks) {
		if len(current.Inputs) > 1 {
			// more than 2 dataset inputs
			break
		}
		if len(current.Inputs) == 0 {
			// no dataset inputs
			break
		}

		if !isMergeableDataset(current.Inputs[0], taskCount) {
			break
		}

		current = current.Inputs[0].Step
		taskCount = len(current.Tasks)

	}
	return current.Id, true
}

// group local steps into one step group
func translateToStepGroups(fc *flow.FlowContext) []*StepGroup {
	// use array instead of map to ensure consistent ordering
	stepId2StepGroup := make([]*StepGroup, len(fc.Steps))
	for _, step := range fc.Steps {
		// println("step:", step.Name, step.Id, "starting...")
		stepId, foundStepId := findAncestorStepId(step)
		if !foundStepId {
			println("step:", step.Id, "Not found stepId.")
			continue
		}
		// println("step:", step.Name, step.Id, "dataset id", stepId)
		if stepId2StepGroup[stepId] == nil {
			stepId2StepGroup[stepId] = NewStepGroup()
			for _, ds := range step.Inputs {
				parentDsId, hasParentIdId := findAncestorStepId(ds.Step)
				if !hasParentIdId {
					// since we add steps following the same order as the code
					log.Panic("parent StepGroup should already be in the map")
				}
				parentSg := stepId2StepGroup[parentDsId]
				if parentSg == nil {
					// since we add steps following the same order as the code
					log.Panic("parent StepGroup should already be in the map")
				}
				stepId2StepGroup[stepId].AddParent(parentSg)
			}
		}
		stepId2StepGroup[stepId].AddStep(step)
	}
	// shrink
	var ret []*StepGroup
	for _, stepGroup := range stepId2StepGroup {
		if stepGroup == nil || len(stepGroup.Steps) == 0 {
			continue
		}
		// println(i, "add step group started by", stepGroup.Steps[0].Name)
		ret = append(ret, stepGroup)
	}
	return ret
}

// group local tasks into one task group
func translateToTaskGroups(stepId2StepGroup []*StepGroup) (ret []*TaskGroup) {
	for _, stepGroup := range stepId2StepGroup {
		assertSameNumberOfTasks(stepGroup.Steps)
		count := len(stepGroup.Steps[0].Tasks)
		// println("dealing with", stepGroup.Steps[0].Name, "tasks:", len(stepGroup.Steps[0].Tasks))
		for i := 0; i < count; i++ {
			tg := NewTaskGroup()
			for _, step := range stepGroup.Steps {
				tg.AddTask(step.Tasks[i])
			}
			// depends on the previous step group
			// MAYBE IMPROVEMENT: depends on a subset of previus shards
			tg.ParentStepGroup = stepGroup
			stepGroup.TaskGroups = append(stepGroup.TaskGroups, tg)
			tg.Id = len(ret)
			ret = append(ret, tg)
		}
	}
	return
}

func assertSameNumberOfTasks(steps []*flow.Step) {
	if len(steps) == 0 {
		return
	}
	count := len(steps[0].Tasks)
	for _, step := range steps {
		if count != len(step.Tasks) {
			log.Fatalf("This should not happen: step %d has %d tasks, but step %d has %d tasks.", steps[0].Id, count, step.Id, len(step.Tasks))
		}
	}
}

func NewStepGroup() *StepGroup {
	return &StepGroup{}
}

func (t *StepGroup) AddStep(Step *flow.Step) *StepGroup {
	t.Steps = append(t.Steps, Step)
	return t
}

func (t *StepGroup) AddParent(parent *StepGroup) *StepGroup {
	t.Parents = append(t.Parents, parent)
	return t
}

func NewTaskGroup() *TaskGroup {
	return &TaskGroup{}
}

func (t *TaskGroup) AddTask(task *flow.Task) *TaskGroup {
	t.Tasks = append(t.Tasks, task)
	return t
}

func (t *TaskGroup) AddParent(parent *TaskGroup) *TaskGroup {
	t.Parents = append(t.Parents, parent)
	return t
}
<commit_msg>better variable names<commit_after>package scheduler

import (
	"log"

	"github.com/chrislusf/glow/flow"
)

type TaskGroupStatus int

const (
	New TaskGroupStatus = iota
	Waiting
	Scheduled
	Completed
	Failed
)

type TaskGroup struct {
	Id              int
	Tasks           []*flow.Task
	Parents         []*TaskGroup
	ParentStepGroup *StepGroup
	Status          TaskGroupStatus
}

type StepGroup struct {
	Steps      []*flow.Step
	Parents    []*StepGroup
	TaskGroups []*TaskGroup
}

func GroupTasks(fc *flow.FlowContext) []*TaskGroup {
	stepGroups := translateToStepGroups(fc)
	return translateToTaskGroups(stepGroups)
}

func isMergeableDataset(ds *flow.Dataset, taskCount int) bool {
	if taskCount != len(ds.Shards) {
		return false
	}
	if taskCount != len(ds.Step.Tasks) {
		return false
	}
	if len(ds.ReadingSteps) > 1 {
		return false
	}
	for _, shard := range ds.Shards {
		if len(shard.ReadingTasks) > 1 {
			return false
		}
	}
	return true
}

// find mergeable parent step or itself if parent is not mergeable
func findAncestorStepId(step *flow.Step) (int, bool) {
	current := step
	taskCount := len(current.Tasks)

	// println("find step", step.Name)

	for taskCount == len(current.Tasks) {
		if len(current.Inputs) > 1 {
			// more than 2 dataset inputs
			break
		}
		if len(current.Inputs) == 0 {
			// no dataset inputs
			break
		}

		if !isMergeableDataset(current.Inputs[0], taskCount) {
			break
		}

		current = current.Inputs[0].Step
		taskCount = len(current.Tasks)

	}
	return current.Id, true
}

// group local steps into one step group
func translateToStepGroups(fc *flow.FlowContext) []*StepGroup {
	// use array instead of map to ensure consistent ordering
	stepId2StepGroup := make([]*StepGroup, len(fc.Steps))
	for _, step := range fc.Steps {
		// println("step:", step.Name, step.Id, "starting...")
		ancestorStepId, foundStepId := findAncestorStepId(step)
		if !foundStepId {
			println("step:", step.Id, "Not found ancestorStepId.")
			continue
		}
		// println("step:", step.Name, step.Id, "ancestorStepId", ancestorStepId)
		if stepId2StepGroup[ancestorStepId] == nil {
			stepId2StepGroup[ancestorStepId] = NewStepGroup()
			for _, ds := range step.Inputs {
				parentDsId, hasParentIdId := findAncestorStepId(ds.Step)
				if !hasParentIdId {
					// since we add steps following the same order as the code
					log.Panic("parent StepGroup should already be in the map")
				}
				parentSg := stepId2StepGroup[parentDsId]
				if parentSg == nil {
					// since we add steps following the same order as the code
					log.Panic("parent StepGroup should already be in the map")
				}
				stepId2StepGroup[ancestorStepId].AddParent(parentSg)
			}
		}
		stepId2StepGroup[ancestorStepId].AddStep(step)
	}
	// shrink
	var ret []*StepGroup
	for _, stepGroup := range stepId2StepGroup {
		if stepGroup == nil || len(stepGroup.Steps) == 0 {
			continue
		}
		// println("add step group started by", stepGroup.Steps[0].Name, "with", len(stepGroup.Steps), "steps")
		ret = append(ret, stepGroup)
	}
	return ret
}

// group local tasks into one task group
func translateToTaskGroups(stepId2StepGroup []*StepGroup) (ret []*TaskGroup) {
	for _, stepGroup := range stepId2StepGroup {
		assertSameNumberOfTasks(stepGroup.Steps)
		count := len(stepGroup.Steps[0].Tasks)
		// println("dealing with", stepGroup.Steps[0].Name, "tasks:", len(stepGroup.Steps[0].Tasks))
		for i := 0; i < count; i++ {
			tg := NewTaskGroup()
			for _, step := range stepGroup.Steps {
				tg.AddTask(step.Tasks[i])
			}
			// depends on the previous step group
			// MAYBE IMPROVEMENT: depends on a subset of previus shards
			tg.ParentStepGroup = stepGroup
			stepGroup.TaskGroups = append(stepGroup.TaskGroups, tg)
			tg.Id = len(ret)
			ret = append(ret, tg)
		}
	}
	return
}

func assertSameNumberOfTasks(steps []*flow.Step) {
	if len(steps) == 0 {
		return
	}
	count := len(steps[0].Tasks)
	for _, step := range steps {
		if count != len(step.Tasks) {
			log.Fatalf("This should not happen: step %d has %d tasks, but step %d has %d tasks.", steps[0].Id, count, step.Id, len(step.Tasks))
		}
	}
}

func NewStepGroup() *StepGroup {
	return &StepGroup{}
}

func (t *StepGroup) AddStep(Step *flow.Step) *StepGroup {
	t.Steps = append(t.Steps, Step)
	return t
}

func (t *StepGroup) AddParent(parent *StepGroup) *StepGroup {
	t.Parents = append(t.Parents, parent)
	return t
}

func NewTaskGroup() *TaskGroup {
	return &TaskGroup{}
}

func (t *TaskGroup) AddTask(task *flow.Task) *TaskGroup {
	t.Tasks = append(t.Tasks, task)
	return t
}

func (t *TaskGroup) AddParent(parent *TaskGroup) *TaskGroup {
	t.Parents = append(t.Parents, parent)
	return t
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/elliptic"
	"encoding/hex"
	"fmt"
	"math/big"
	"strconv"
	"strings"
	"testing"
)

func TestPath(t *testing.T) {
	//key for tag
	D := new(big.Int)
	D, _ = new(big.Int).SetString("399681106706823979931786798252509423226869972"+
		"6722344370689730210711314983767775860556101498400185744208447673206609026"+
		"128894016152514163591905578729891874833", 10)
	privT := PrivteKeyFromD(*D)
	keyT := elliptic.Marshal(privT.PublicKey.Curve,
		privT.PublicKey.X, privT.PublicKey.Y)
	hkidT, _ := hex.DecodeString(GenerateHKID(privT)) //gen HKID for tag
	err := PostKey(privT)                             //place key for tag
	if err != nil {
		fmt.Println(err)
	}
	err = PostBlob(keyT) //store public key
	if err != nil {
		fmt.Println(err)
	}

	//key for commit
	//r, ok := new(big.Int).SetString(s, 16)
	D, _ = new(big.Int).SetString("462981482389329648001641133480879383618612455"+
		"9723200979962176753724976464088706463001383556112424820911870650421151988"+
		"906751710824965155500230480521264034469", 10)
	privC := PrivteKeyFromD(*D)
	keyC := elliptic.Marshal(privC.PublicKey.Curve,
		privC.PublicKey.X, privC.PublicKey.Y)
	hkidC, _ := hex.DecodeString(GenerateHKID(privC)) //gen HKID for commit
	err = PostKey(privC)                              //place key for commit
	if err != nil {
		fmt.Println(err)
	}
	err = PostBlob(keyC) //store public key
	if err != nil {
		fmt.Println(err)
	}

	//Post blob
	testBlob := blob([]byte("testing")) //gen test blob
	err = PostBlob(testBlob)            //store test blob
	if err != nil {
		fmt.Println(err)
	}

	//post tag
	testTagPointingToTestBlob := NewTag(testBlob.Hash(),
		"blob",
		"testBlob",
		hkidT) //gen test tag
	err = PostTag(testTagPointingToTestBlob) //post test tag
	if err != nil {
		fmt.Println(err)
	}

	//post list
	testListPiontingToTestTag := NewList(testTagPointingToTestBlob.Hkid(),
		"tag",
		"testTag") //gen test list
	err = PostBlob(testListPiontingToTestTag.Bytes()) //store test list
	if err != nil {
		fmt.Println(err)
	}

	// post commit
	testCommitPointingToTestList := NewCommit(testListPiontingToTestTag.Hash(),
		hkidC) //gen test commit
	err = PostCommit(testCommitPointingToTestList) //post test commit
	if err != nil {
		fmt.Println(err)
	}

	//print
	fmt.Printf("%s\n%s\n%s\n%s\n",
		hex.EncodeToString(testBlob.Hash()),
		hex.EncodeToString(testTagPointingToTestBlob.Hash()),
		hex.EncodeToString(testListPiontingToTestTag.Hash()),
		hex.EncodeToString(testCommitPointingToTestList.Hash()))

	//get commit
	hkid, _ := hex.DecodeString("1312ac161875b270da2ae4e1471ba94a" +
		"9883419250caa4c2f1fd80a91b37907e")
	commitbytes, err := GetCommit(hkid)
	if err != nil {
		panic(err)
	}
	commitStrings := strings.Split(string(commitbytes), ",\n")
	listHash, _ := hex.DecodeString(commitStrings[0])
	version, _ := strconv.ParseInt(commitStrings[1], 10, 64)
	chkid, _ := hex.DecodeString(commitStrings[2])
	signature, _ := hex.DecodeString(commitStrings[3])
	testcommit := commit{listHash, version, chkid, signature}
	//fmt.Println(testcommit)
	fmt.Printf("authentic commit:%v\n", testcommit.Verifiy())

	//get list
	listbytes, err := GetBlob(listHash)
	if err != nil {
		panic(err)
	}
	listEntries := strings.Split(string(listbytes), "\n")
	entries := []entry{}
	cols := []string{}
	for _, element := range listEntries {
		cols = strings.Split(element, ",")
		entryHash, _ := hex.DecodeString(cols[0])
		entryTypeString := cols[1]
		entryNameSegment := cols[2]
		entries = append(entries, entry{entryHash, entryTypeString, entryNameSegment})
	}
	testlist := list(entries)
	fmt.Printf("authentic list:%v\n", bytes.Equal(listHash, testlist.Hash()))
	//get tag
	tagBytes, err := GetTag(testlist[0].Hash, "testBlob")
	tagStrings := strings.Split(string(tagBytes), ",\n")
	tagHashBytes, _ := hex.DecodeString(tagStrings[0])
	tagTypeString := tagStrings[1]
	tagNameSegment := tagStrings[2]
	tagVersion, _ := strconv.ParseInt(tagStrings[3], 10, 64)
	tagHkid, _ := hex.DecodeString(tagStrings[4])
	tagSignature, _ := hex.DecodeString(tagStrings[5])
	testTag := Tag{tagHashBytes, tagTypeString, tagNameSegment, tagVersion, tagHkid,
		tagSignature}
	fmt.Printf("authentic tag:%v\n", testTag.Verifiy())
	//get blob
	blobBytes, _ := GetBlob(tagHashBytes)
	testBlob = blob(blobBytes)
	fmt.Printf("authentic blob:%v\n", bytes.Equal(tagHashBytes, testBlob.Hash()))
}

/*func TestNewCommit(t *testing.T) {
	listHash := []byte{0xfa,
		0x84, 0xff, 0xaa, 0xe4, 0xd6, 0x5f, 0x49, 0x67, 0x67, 0x8c, 0x95, 0xb7, 0xf9, 0x6d, 0x61,
		0xe0, 0x67, 0x2b, 0x77, 0xe2, 0x67, 0x26, 0x78, 0x44, 0x44, 0x95, 0x24, 0x55, 0x07, 0x56, 0xb4}
	hkid := []byte{0xfa,
		0x84, 0xff, 0xaa, 0xe4, 0xd6, 0x5f, 0x49, 0x67, 0x67, 0x8c, 0x95, 0xb7, 0xf9, 0x6d, 0x61,
		0xe0, 0x67, 0x2b, 0x77, 0xe2, 0x67, 0x26, 0x78, 0x44, 0x44, 0x95, 0x24, 0x55, 0x07, 0x56, 0xb4}
	fmt.Print(listHash, hkid)
	c := NewCommit(listHash, hkid)
	fmt.Print(c.Verifiy())
}*/

/*
func TestKeyGen(t *testing.T) {
	c := elliptic.P521()
	priv, err := ecdsa.GenerateKey(c, rand.Reader)
	if err != nil {
		t.Errorf("Error %v", err)
	}
	fmt.Printf("TestKeyGen\nX = %v\nY = %v\nD = %v\n", priv.PublicKey.X, priv.PublicKey.Y, priv.D)
	//priv.PublicKey.X, priv.PublicKey.Y = c.ScalarBaseMult(priv.D.Bytes())
	//fmt.Printf("X = %v\nY = %v\nD = %v\n", priv.PublicKey.X, priv.PublicKey.Y, priv.D)
	test(priv, t)
	//SaveKey(priv)
	SavePrivateKey(priv)
	priv, err = LoadPrivateKey(GenerateHKID(priv))
	fmt.Printf("\nhkid:%v\n", GenerateHKID(priv))
}
*/

/*
func TestCommitgen(t *testing.T) {
	priv, _ := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
	SavePrivateKey(priv)
	//fmt.Printf("\nTestCommitgen\nX = %v\nY = %v\nD = %v\n", priv.PublicKey.X, priv.PublicKey.Y, priv.D)
	hkid := GenerateHKID(priv)
	reconstructedkey, err := LoadPrivateKey(hkid)
	if err != nil {
		panic(err)
	}
	reconstructedhkid := GenerateHKID(reconstructedkey)
	fmt.Printf("\n%v\n%v", hkid, reconstructedhkid)
	fmt.Printf("\nTestCommitgen\n%v", GenerateCommit([]byte("testing"), priv))
}
*/
/*
func TestTaggen(t *testing.T) {
	priv, _ := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
	//test(priv)
	fmt.Printf("\nTestTaggen\n%v\n", GenerateTag([]byte("testing"), "blob", "test", priv))
}
*/
/*
func test(priv *ecdsa.PrivateKey, t *testing.T) {
	mar := elliptic.Marshal(
		priv.PublicKey.Curve,
		priv.PublicKey.X,
		priv.PublicKey.Y)
	x, y := elliptic.Unmarshal(elliptic.P521(), mar)
	maredPublicKey := new(ecdsa.PublicKey)
	maredPublicKey.Curve = elliptic.P521()
	maredPublicKey.X = x
	maredPublicKey.Y = y

	hashed := []byte("testing")
	r, s, _ := ecdsa.Sign(rand.Reader, priv, hashed)
	valid := ecdsa.Verify(maredPublicKey, hashed, r, s)
	if valid != true {
		t.Errorf("failed Verify\n")
	}
	invalid := ecdsa.Verify(maredPublicKey, []byte("fail testing"), r, s)
	if invalid != false {
		t.Errorf("failed falsify\n")
	}
	//fmt.Printf("valid:%v in:%v marsize:%v bits\n\n\n", valid, invalid, len(mar)*8)
}
*/
<commit_msg>re-enabled key gen test<commit_after>package main

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"math/big"
	"strconv"
	"strings"
	"testing"
)

func TestPath(t *testing.T) {
	//key for tag
	D := new(big.Int)
	D, _ = new(big.Int).SetString("399681106706823979931786798252509423226869972"+
		"6722344370689730210711314983767775860556101498400185744208447673206609026"+
		"128894016152514163591905578729891874833", 10)
	privT := PrivteKeyFromD(*D)
	keyT := elliptic.Marshal(privT.PublicKey.Curve,
		privT.PublicKey.X, privT.PublicKey.Y)
	hkidT, _ := hex.DecodeString(GenerateHKID(privT)) //gen HKID for tag
	err := PostKey(privT)                             //place key for tag
	if err != nil {
		fmt.Println(err)
	}
	err = PostBlob(keyT) //store public key
	if err != nil {
		fmt.Println(err)
	}

	//key for commit
	//r, ok := new(big.Int).SetString(s, 16)
	D, _ = new(big.Int).SetString("462981482389329648001641133480879383618612455"+
		"9723200979962176753724976464088706463001383556112424820911870650421151988"+
		"906751710824965155500230480521264034469", 10)
	privC := PrivteKeyFromD(*D)
	keyC := elliptic.Marshal(privC.PublicKey.Curve,
		privC.PublicKey.X, privC.PublicKey.Y)
	hkidC, _ := hex.DecodeString(GenerateHKID(privC)) //gen HKID for commit
	err = PostKey(privC)                              //place key for commit
	if err != nil {
		fmt.Println(err)
	}
	err = PostBlob(keyC) //store public key
	if err != nil {
		fmt.Println(err)
	}

	//Post blob
	testBlob := blob([]byte("testing")) //gen test blob
	err = PostBlob(testBlob)            //store test blob
	if err != nil {
		fmt.Println(err)
	}

	//post tag
	testTagPointingToTestBlob := NewTag(testBlob.Hash(),
		"blob",
		"testBlob",
		hkidT) //gen test tag
	err = PostTag(testTagPointingToTestBlob) //post test tag
	if err != nil {
		fmt.Println(err)
	}

	//post list
	testListPiontingToTestTag := NewList(testTagPointingToTestBlob.Hkid(),
		"tag",
		"testTag") //gen test list
	err = PostBlob(testListPiontingToTestTag.Bytes()) //store test list
	if err != nil {
		fmt.Println(err)
	}

	// post commit
	testCommitPointingToTestList := NewCommit(testListPiontingToTestTag.Hash(),
		hkidC) //gen test commit
	err = PostCommit(testCommitPointingToTestList) //post test commit
	if err != nil {
		fmt.Println(err)
	}

	//print
	fmt.Printf("%s\n%s\n%s\n%s\n",
		hex.EncodeToString(testBlob.Hash()),
		hex.EncodeToString(testTagPointingToTestBlob.Hash()),
		hex.EncodeToString(testListPiontingToTestTag.Hash()),
		hex.EncodeToString(testCommitPointingToTestList.Hash()))

	//get commit
	hkid, _ := hex.DecodeString("1312ac161875b270da2ae4e1471ba94a" +
		"9883419250caa4c2f1fd80a91b37907e")
	commitbytes, err := GetCommit(hkid)
	if err != nil {
		panic(err)
	}
	commitStrings := strings.Split(string(commitbytes), ",\n")
	listHash, _ := hex.DecodeString(commitStrings[0])
	version, _ := strconv.ParseInt(commitStrings[1], 10, 64)
	chkid, _ := hex.DecodeString(commitStrings[2])
	signature, _ := hex.DecodeString(commitStrings[3])
	testcommit := commit{listHash, version, chkid, signature}
	//fmt.Println(testcommit)
	fmt.Printf("authentic commit:%v\n", testcommit.Verifiy())

	//get list
	listbytes, err := GetBlob(listHash)
	if err != nil {
		panic(err)
	}
	listEntries := strings.Split(string(listbytes), "\n")
	entries := []entry{}
	cols := []string{}
	for _, element := range listEntries {
		cols = strings.Split(element, ",")
		entryHash, _ := hex.DecodeString(cols[0])
		entryTypeString := cols[1]
		entryNameSegment := cols[2]
		entries = append(entries, entry{entryHash, entryTypeString, entryNameSegment})
	}
	testlist := list(entries)
	fmt.Printf("authentic list:%v\n", bytes.Equal(listHash, testlist.Hash()))
	//get tag
	tagBytes, err := GetTag(testlist[0].Hash, "testBlob")
	tagStrings := strings.Split(string(tagBytes), ",\n")
	tagHashBytes, _ := hex.DecodeString(tagStrings[0])
	tagTypeString := tagStrings[1]
	tagNameSegment := tagStrings[2]
	tagVersion, _ := strconv.ParseInt(tagStrings[3], 10, 64)
	tagHkid, _ := hex.DecodeString(tagStrings[4])
	tagSignature, _ := hex.DecodeString(tagStrings[5])
	testTag := Tag{tagHashBytes, tagTypeString, tagNameSegment, tagVersion, tagHkid,
		tagSignature}
	fmt.Printf("authentic tag:%v\n", testTag.Verifiy())
	//get blob
	blobBytes, _ := GetBlob(tagHashBytes)
	testBlob = blob(blobBytes)
	fmt.Printf("authentic blob:%v\n", bytes.Equal(tagHashBytes, testBlob.Hash()))
}

func TestKeyGen(t *testing.T) {
	c := elliptic.P521()
	priv, err := ecdsa.GenerateKey(c, rand.Reader)
	if err != nil {
		t.Errorf("Error %v", err)
	}
	fmt.Printf("TestKeyGen\nX = %v\nY = %v\nD = %v\n", priv.PublicKey.X, priv.PublicKey.Y, priv.D)
	err = PostKey(priv)
	if err != nil {
		t.Errorf("Error %v", err)
	}
	PostBlob(elliptic.Marshal(priv.PublicKey.Curve,
		priv.PublicKey.X, priv.PublicKey.Y))
}

/*
func TestCommitgen(t *testing.T) {
	priv, _ := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
	SavePrivateKey(priv)
	//fmt.Printf("\nTestCommitgen\nX = %v\nY = %v\nD = %v\n", priv.PublicKey.X, priv.PublicKey.Y, priv.D)
	hkid := GenerateHKID(priv)
	reconstructedkey, err := LoadPrivateKey(hkid)
	if err != nil {
		panic(err)
	}
	reconstructedhkid := GenerateHKID(reconstructedkey)
	fmt.Printf("\n%v\n%v", hkid, reconstructedhkid)
	fmt.Printf("\nTestCommitgen\n%v", GenerateCommit([]byte("testing"), priv))
}
*/
/*
func TestTaggen(t *testing.T) {
	priv, _ := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
	//test(priv)
	fmt.Printf("\nTestTaggen\n%v\n", GenerateTag([]byte("testing"), "blob", "test", priv))
}
*/
/*
func test(priv *ecdsa.PrivateKey, t *testing.T) {
	mar := elliptic.Marshal(
		priv.PublicKey.Curve,
		priv.PublicKey.X,
		priv.PublicKey.Y)
	x, y := elliptic.Unmarshal(elliptic.P521(), mar)
	maredPublicKey := new(ecdsa.PublicKey)
	maredPublicKey.Curve = elliptic.P521()
	maredPublicKey.X = x
	maredPublicKey.Y = y

	hashed := []byte("testing")
	r, s, _ := ecdsa.Sign(rand.Reader, priv, hashed)
	valid := ecdsa.Verify(maredPublicKey, hashed, r, s)
	if valid != true {
		t.Errorf("failed Verify\n")
	}
	invalid := ecdsa.Verify(maredPublicKey, []byte("fail testing"), r, s)
	if invalid != false {
		t.Errorf("failed falsify\n")
	}
	//fmt.Printf("valid:%v in:%v marsize:%v bits\n\n\n", valid, invalid, len(mar)*8)
}
*/
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes 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 logs

import (
	"flag"
	"log"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/util/wait"
)

var logFlushFreq = pflag.Duration("log-flush-frequency", 5*time.Second, "Maximum number of seconds between log flushes")

// TODO(thockin): This is temporary until we agree on log dirs and put those into each cmd.
func init() {
	flag.Set("logtostderr", "true")
}

// GlogWriter serves as a bridge between the standard log package and the glog package.
type GlogWriter struct{}

// Write implements the io.Writer interface.
func (writer GlogWriter) Write(data []byte) (n int, err error) {
	glog.Info(string(data))
	return len(data), nil
}

// InitLogs initializes logs the way we want for kubernetes.
func InitLogs() {
	log.SetOutput(GlogWriter{})
	log.SetFlags(0)
	// The default glog flush interval is 30 seconds, which is frighteningly long.
	go wait.Until(glog.Flush, *logFlushFreq, wait.NeverStop)
}

// FlushLogs flushes logs immediately.
func FlushLogs() {
	glog.Flush()
}

// NewLogger creates a new log.Logger which sends logs to glog.Info.
func NewLogger(prefix string) *log.Logger {
	return log.New(GlogWriter{}, prefix, 0)
}
<commit_msg>fix comment typo and use wait.Forever<commit_after>/*
Copyright 2014 The Kubernetes 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 logs

import (
	"flag"
	"log"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/util/wait"
)

var logFlushFreq = pflag.Duration("log-flush-frequency", 5*time.Second, "Maximum number of seconds between log flushes")

// TODO(thockin): This is temporary until we agree on log dirs and put those into each cmd.
func init() {
	flag.Set("logtostderr", "true")
}

// GlogWriter serves as a bridge between the standard log package and the glog package.
type GlogWriter struct{}

// Write implements the io.Writer interface.
func (writer GlogWriter) Write(data []byte) (n int, err error) {
	glog.Info(string(data))
	return len(data), nil
}

// InitLogs initializes logs the way we want for kubernetes.
func InitLogs() {
	log.SetOutput(GlogWriter{})
	log.SetFlags(0)
	// The default glog flush interval is 5 seconds.
	go wait.Forever(glog.Flush, *logFlushFreq)
}

// FlushLogs flushes logs immediately.
func FlushLogs() {
	glog.Flush()
}

// NewLogger creates a new log.Logger which sends logs to glog.Info.
func NewLogger(prefix string) *log.Logger {
	return log.New(GlogWriter{}, prefix, 0)
}
<|endoftext|>


<commit_before>package dpn

/*
restobjects.go includes a number of structures used in communicating
with the DPN REST service. This file defines only those DPN REST structures
related to storing bags and creating and fulfilling replication requests.
*/

import (
	"time"
)

// DPNFixity represents a checksum for a bag in the DPN REST
// service.
type DPNFixity struct {

	// The algorithm used to check the fixity. Usually 'sha256',
	// but others may be valid in the future.
	Algorithm string

	// The fixity digest, as a hex-encoded string.
	Digest string

	// The datetime at which this digest was calculated.
	CreatedAt time.Time
}

// DPNBag represents a Bag object in the DPN REST service.
// Like all of the DPN REST objects, it contains metadata only.
type DPNBag struct {

	// UUID is the unique identifier for a bag
	UUID               string

	// LocalId is the depositor's local identifier for a bag.
	LocalId            string

	// Size is the size, in bytes of the bag.
	Size               uint64

	// FirstVersionUUID is the UUID of the first version
	// of this bag.
	FirstVersionUUID   string

	// Version is the version or revision number of the bag. Starts at 1.
	Version            uint32

	// OriginalNode is the node that first ingested or produced the bag.
	OriginalNode       string

	// BagType is one of 'D' (Data), 'R' (Rights) or 'I' (Interpretive)
	BagType            string

	// Rights is a list of UUIDs of rights objects for this bag.
	Rights             []string

	// Interpretive is a list of UUIDs of interpretive objects for this bag.
	Interpretive       []string

	// AdminNode is the authoritative node for this bag. If various nodes
	// have conflicting registry info for this bag, the admin node wins.
	// The admin node also has some authority in restoring and (if its ever
	// possible) deleting bags.
	AdminNode          string
}

type DPNReplicationTransfer struct {

}

type DPNRestoreTransfer struct {

}
<commit_msg>Added preliminary DPN REST objects<commit_after>package dpn

/*
restobjects.go includes a number of structures used in communicating
with the DPN REST service. This file defines only those DPN REST structures
related to storing bags and creating and fulfilling replication requests.

More info about the DPN REST service is available at:

https://github.com/dpn-admin/DPN-REST-Wiki
https://github.com/dpn-admin/DPN-REST
*/

import (
	"time"
)

// DPNFixity represents a checksum for a bag in the DPN REST
// service.
type DPNFixity struct {

	// The algorithm used to check the fixity. Usually 'sha256',
	// but others may be valid in the future.
	Algorithm string

	// The fixity digest, as a hex-encoded string.
	Digest string

	// The datetime at which this digest was calculated.
	CreatedAt time.Time
}

// DPNBag represents a Bag object in the DPN REST service.
// Like all of the DPN REST objects, it contains metadata only.
type DPNBag struct {

	// UUID is the unique identifier for a bag
	UUID               string               `json:"uuid"`

	// LocalId is the depositor's local identifier for a bag.
	LocalId            string               `json:"local_id"`

	// Size is the size, in bytes of the bag.
	Size               uint64               `json:"size"`

	// FirstVersionUUID is the UUID of the first version
	// of this bag.
	FirstVersionUUID   string               `json:"first_version_uuid"`

	// Version is the version or revision number of the bag. Starts at 1.
	Version            uint32               `json:"version"`

	// IngestNode is the node that first ingested or produced the bag.
	IngestNode       string                 `json:"ingest_node"`

	// AdminNode is the authoritative node for this bag. If various nodes
	// have conflicting registry info for this bag, the admin node wins.
	// The admin node also has some authority in restoring and (if its ever
	// possible) deleting bags.
	AdminNode          string               `json:"ingest_node"`

	// BagType is one of 'D' (Data), 'R' (Rights) or 'I' (Interpretive)
	BagType            string               `json:"bag_type"`

	// Rights is a list of UUIDs of rights objects for this bag.
	Rights             []string             `json:"rights"`

	// Interpretive is a list of UUIDs of interpretive objects for this bag.
	Interpretive       []string             `json:"interpretive"`

	// ReplicatingNodes is a list of one more nodes that has stored
	// copies of this bag. The items in the list are node namespaces,
	// which are strings. E.g. ['aptrust', 'chron', 'tdr']
	ReplicatingNodes   []string             `json:"replicating_nodes"`

	// Fixities is a list of hashes. In each hash, the key is the
	// fixity algorithm and the value is the fixity digest. So you
	// might see something like this:
	// [{ "sha256": "1fdc62a", "sha512": "9f8d23ae" }]
	Fixities           []map[string]string  `json:"fixities"`

	// CreatedAt is when this record was created.
	CreatedAt          time.Time            `json:"created_at"`

	// UpdatedAt is when this record was last updated.
	UpdatedAt          time.Time            `json:"updated_at"`
}

type DPNReplicationTransfer struct {
	FromNode        string       `json:"from_node"`
	ToNode          string       `json:"to_node"`
	UUID            string       `json:"bag"`
	ReplicationId   string       `json:"replication_id"`
	FixityAlgorithm string       `json:"fixity_algorithm"`
	FixityNonce     string       `json:"fixity_nonce"`
	FixityValue     string       `json:"fixity_value"`
	FixityAccept    bool         `json:"fixity_accept"`
	BagValid        bool         `json:"bag_valid"`
	Status          string       `json:"status"`
	Protocol        string       `json:"protocol"`
	Link            string       `json:"link"`
	CreatedAt       time.Time    `json:"created_at"`
	UpdatedAt       time.Time    `json:"updated_at"`
}

type DPNRestoreTransfer struct {
	FromNode        string       `json:"from_node"`
	ToNode          string       `json:"to_node"`
	UUID            string       `json:"bag"`
	Status          string       `json:"status"`
	Protocol        string       `json:"protocol"`
	Link            string       `json:"link"`
	CreatedAt       time.Time    `json:"created_at"`
	UpdatedAt       time.Time    `json:"updated_at"`
}
<|endoftext|>


<commit_before>package ddtxn

import (
	"flag"
	"log"
	"runtime/debug"
)

var SampleRate = flag.Int64("sr", 1000, "Sample every sr nanoseconds\n")

// TODO: 2PL

type Write struct {
	key    Key
	br     *BRecord
	v      Value
	op     KeyType
	create bool
	locked bool
	vint32 int32

	// TODO: Handle writing more than once to a list per txn
	ve Entry
	dd bool
}

// Not threadsafe.  Tracks execution of transaction.
type ETransaction struct {
	read       []*BRecord
	lasts      []uint64
	w          *Worker
	s          *Store
	ls         *LocalStore
	writes     []Write
	t          int64 // Used just as a rough count
	any_marked bool
	padding    [128]byte
}

func StartTransaction(w *Worker) *ETransaction {
	tx := &ETransaction{
		read:   make([]*BRecord, 0, 100),
		lasts:  make([]uint64, 0, 100),
		writes: make([]Write, 0, 100),
		w:      w,
		s:      w.store,
		ls:     w.local_store,
	}
	return tx
}

func (tx *ETransaction) Reset() {
	tx.lasts = tx.lasts[:0]
	tx.read = tx.read[:0]
	tx.writes = tx.writes[:0]
	tx.ls.count = (*SysType == DOPPEL && tx.t%*SampleRate == 0)
	if tx.ls.count {
		tx.w.Nstats[NSAMPLES]++
	}
	tx.t++
}

func (tx *ETransaction) Read(k Key) (*BRecord, error) {
	if *SysType == DOPPEL {
		if tx.ls.phase == SPLIT {
			if tx.any_marked && tx.s.IsDD(k) {
				if tx.ls.count {
					tx.ls.candidates.Stash(k)
				}
				return nil, ESTASH
			}
		}
	}
	// TODO: If I wrote the key, return that value instead
	br, err := tx.s.getKey(k)
	//tx.w.Nstats[NGETKEYCALLS]++
	if *CountKeys {
		p, r := UndoCKey(k)
		if r == 117 {
			tx.w.NKeyAccesses[p]++
		}
	}
	if err != nil {
		return nil, err
	}
	ok, last := br.IsUnlocked()
	// if locked and not by me, abort
	// else note the last timestamp, save it, return value
	if !ok {
		if tx.ls.count {
			tx.ls.candidates.Conflict(k)
		}
		tx.Abort()
		return nil, EABORT
	}
	n := len(tx.read)
	tx.read = tx.read[0 : n+1]
	tx.read[n] = br
	tx.lasts = tx.lasts[0 : n+1]
	tx.lasts[n] = last
	return br, nil
}

func (tx *ETransaction) add(k Key, v Value, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].v = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) addInt32(k Key, v int32, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].vint32 = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) addList(k Key, v Entry, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].ve = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) WriteInt32(k Key, a int32, op KeyType) {
	tx.addInt32(k, a, op, false)
}

func (tx *ETransaction) Write(k Key, v Value, kt KeyType) {
	if kt == SUM || kt == MAX {
		tx.addInt32(k, v.(int32), kt, true)
		return
	}
	tx.add(k, v, kt, true)
}

func (tx *ETransaction) WriteList(k Key, l Entry, kt KeyType) {
	if kt != LIST {
		log.Fatalf("Not a list\n")
	}
	tx.addList(k, l, kt, true)
}

func (tx *ETransaction) Abort() TID {
	for i, _ := range tx.writes {
		if tx.writes[i].locked {
			tx.writes[i].br.Unlock(0)
		}
	}
	return 0
}

func (tx *ETransaction) Commit() TID {
	// for each write key
	//  if global get from global store and lock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if *SysType == DOPPEL && tx.ls.phase == SPLIT && tx.s.IsDD(w.key) {
			w.dd = true
			continue
		}
		if w.br == nil {
			br, err := tx.s.getKey(w.key)
			//tx.w.Nstats[NGETKEYCALLS]++
			if *CountKeys {
				p, r := UndoCKey(w.key)
				if r == 117 {
					tx.w.NKeyAccesses[p]++
				}
			}
			if br == nil || err != nil {
				switch w.op {
				case SUM:
					br = tx.s.CreateInt32Key(w.key, 0, w.op)
				case MAX:
					br = tx.s.CreateInt32Key(w.key, 0, w.op)
				case LIST:
					br = tx.s.CreateKey(w.key, nil, LIST)
				default:
					if w.v == nil {
						br = tx.s.CreateKey(w.key, "", WRITE)
					} else {
						br = tx.s.CreateKey(w.key, w.v, WRITE)
					}
				}
			}
			w.br = br
		}
		if !w.br.Lock() {
			if tx.ls.count {
				tx.ls.candidates.Conflict(w.key)
			}
			return tx.Abort()
		}
		w.locked = true
	}
	// TODO: acquire timestamp higher than anything i've read or am
	// writing
	tid := tx.w.commitTID()

	// for each read key
	//  verify
	if len(tx.read) != len(tx.lasts) {
		debug.PrintStack()
		log.Fatalf("Mismatch in lengths reads: %v, lasts: %v\n", tx.read, tx.lasts)
	}
	for i, _ := range tx.read {
		// Would have checked for dd earlier
		if tx.ls.count {
			tx.ls.candidates.Read(tx.read[i].key)
		}
		rd := false
		if !tx.read[i].Verify(tx.lasts[i]) {
			for j, _ := range tx.writes {
				if tx.writes[j].key == tx.read[i].key {
					// We would have aborted if we did not successfully
					// lock this earlier
					rd = true
					break
				}
			}
			if rd {
				continue
			}
			if tx.ls.count {
				tx.ls.candidates.Conflict(tx.read[i].key)
			}
			return tx.Abort()
		}
	}
	// for each write key
	//  if dd and split phase, apply locally
	//  else apply globally and unlock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if tx.ls.phase == SPLIT && w.dd {
			if tx.ls.count {
				tx.ls.candidates.Write(w.key)
			}
			//tx.w.Nstats[NDDWRITES]++
			switch w.op {
			case SUM:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			case MAX:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			case LIST:
				tx.ls.ApplyList(w.key, w.ve)
			default:
				tx.ls.Apply(w.key, w.op, w.v, w.op)
			}
		} else {
			switch w.op {
			case SUM:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			case MAX:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			case LIST:
				tx.s.Set(w.br, w.ve, w.op)
			default:
				tx.s.Set(w.br, w.v, w.op)
			}
			w.br.Unlock(tid)
		}
	}
	return tid
}
<commit_msg>use any marked<commit_after>package ddtxn

import (
	"flag"
	"log"
	"runtime/debug"
)

var SampleRate = flag.Int64("sr", 1000, "Sample every sr nanoseconds\n")

// TODO: 2PL

type Write struct {
	key    Key
	br     *BRecord
	v      Value
	op     KeyType
	create bool
	locked bool
	vint32 int32

	// TODO: Handle writing more than once to a list per txn
	ve Entry
	dd bool
}

// Not threadsafe.  Tracks execution of transaction.
type ETransaction struct {
	read       []*BRecord
	lasts      []uint64
	w          *Worker
	s          *Store
	ls         *LocalStore
	writes     []Write
	t          int64 // Used just as a rough count
	any_marked bool
	padding    [128]byte
}

func StartTransaction(w *Worker) *ETransaction {
	tx := &ETransaction{
		read:   make([]*BRecord, 0, 100),
		lasts:  make([]uint64, 0, 100),
		writes: make([]Write, 0, 100),
		w:      w,
		s:      w.store,
		ls:     w.local_store,
	}
	return tx
}

func (tx *ETransaction) Reset() {
	tx.lasts = tx.lasts[:0]
	tx.read = tx.read[:0]
	tx.writes = tx.writes[:0]
	tx.ls.count = (*SysType == DOPPEL && tx.t%*SampleRate == 0)
	if tx.ls.count {
		tx.w.Nstats[NSAMPLES]++
	}
	tx.t++
}

func (tx *ETransaction) Read(k Key) (*BRecord, error) {
	if *SysType == DOPPEL {
		if tx.ls.phase == SPLIT {
			if tx.any_marked && tx.s.IsDD(k) {
				if tx.ls.count {
					tx.ls.candidates.Stash(k)
				}
				return nil, ESTASH
			}
		}
	}
	// TODO: If I wrote the key, return that value instead
	br, err := tx.s.getKey(k)
	//tx.w.Nstats[NGETKEYCALLS]++
	if *CountKeys {
		p, r := UndoCKey(k)
		if r == 117 {
			tx.w.NKeyAccesses[p]++
		}
	}
	if err != nil {
		return nil, err
	}
	ok, last := br.IsUnlocked()
	// if locked and not by me, abort
	// else note the last timestamp, save it, return value
	if !ok {
		if tx.ls.count {
			tx.ls.candidates.Conflict(k)
		}
		tx.Abort()
		return nil, EABORT
	}
	n := len(tx.read)
	tx.read = tx.read[0 : n+1]
	tx.read[n] = br
	tx.lasts = tx.lasts[0 : n+1]
	tx.lasts[n] = last
	return br, nil
}

func (tx *ETransaction) add(k Key, v Value, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].v = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) addInt32(k Key, v int32, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].vint32 = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) addList(k Key, v Entry, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].ve = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) WriteInt32(k Key, a int32, op KeyType) {
	tx.addInt32(k, a, op, false)
}

func (tx *ETransaction) Write(k Key, v Value, kt KeyType) {
	if kt == SUM || kt == MAX {
		tx.addInt32(k, v.(int32), kt, true)
		return
	}
	tx.add(k, v, kt, true)
}

func (tx *ETransaction) WriteList(k Key, l Entry, kt KeyType) {
	if kt != LIST {
		log.Fatalf("Not a list\n")
	}
	tx.addList(k, l, kt, true)
}

func (tx *ETransaction) Abort() TID {
	for i, _ := range tx.writes {
		if tx.writes[i].locked {
			tx.writes[i].br.Unlock(0)
		}
	}
	return 0
}

func (tx *ETransaction) Commit() TID {
	// for each write key
	//  if global get from global store and lock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if *SysType == DOPPEL && tx.ls.phase == SPLIT {
			if tx.any_marked && tx.s.IsDD(w.key) {
				w.dd = true
				continue
			}
		}
		if w.br == nil {
			br, err := tx.s.getKey(w.key)
			//tx.w.Nstats[NGETKEYCALLS]++
			if *CountKeys {
				p, r := UndoCKey(w.key)
				if r == 117 {
					tx.w.NKeyAccesses[p]++
				}
			}
			if br == nil || err != nil {
				switch w.op {
				case SUM:
					br = tx.s.CreateInt32Key(w.key, 0, w.op)
				case MAX:
					br = tx.s.CreateInt32Key(w.key, 0, w.op)
				case LIST:
					br = tx.s.CreateKey(w.key, nil, LIST)
				default:
					if w.v == nil {
						br = tx.s.CreateKey(w.key, "", WRITE)
					} else {
						br = tx.s.CreateKey(w.key, w.v, WRITE)
					}
				}
			}
			w.br = br
		}
		if !w.br.Lock() {
			if tx.ls.count {
				tx.ls.candidates.Conflict(w.key)
			}
			return tx.Abort()
		}
		w.locked = true
	}
	// TODO: acquire timestamp higher than anything i've read or am
	// writing
	tid := tx.w.commitTID()

	// for each read key
	//  verify
	if len(tx.read) != len(tx.lasts) {
		debug.PrintStack()
		log.Fatalf("Mismatch in lengths reads: %v, lasts: %v\n", tx.read, tx.lasts)
	}
	for i, _ := range tx.read {
		// Would have checked for dd earlier
		if tx.ls.count {
			tx.ls.candidates.Read(tx.read[i].key)
		}
		rd := false
		if !tx.read[i].Verify(tx.lasts[i]) {
			for j, _ := range tx.writes {
				if tx.writes[j].key == tx.read[i].key {
					// We would have aborted if we did not successfully
					// lock this earlier
					rd = true
					break
				}
			}
			if rd {
				continue
			}
			if tx.ls.count {
				tx.ls.candidates.Conflict(tx.read[i].key)
			}
			return tx.Abort()
		}
	}
	// for each write key
	//  if dd and split phase, apply locally
	//  else apply globally and unlock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if tx.ls.phase == SPLIT && w.dd {
			if tx.ls.count {
				tx.ls.candidates.Write(w.key)
			}
			//tx.w.Nstats[NDDWRITES]++
			switch w.op {
			case SUM:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			case MAX:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			case LIST:
				tx.ls.ApplyList(w.key, w.ve)
			default:
				tx.ls.Apply(w.key, w.op, w.v, w.op)
			}
		} else {
			switch w.op {
			case SUM:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			case MAX:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			case LIST:
				tx.s.Set(w.br, w.ve, w.op)
			default:
				tx.s.Set(w.br, w.v, w.op)
			}
			w.br.Unlock(tid)
		}
	}
	return tid
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 incrementaldpf calls the C++ incremental DPF functions from Go.
package incrementaldpf

// #cgo CFLAGS: -g -Wall
// #include <stdbool.h>
// #include <stdlib.h>
// #include "pipeline/distributed_point_function_c_bridge.h"
import (
	"C"
)

import (
	"errors"
	"fmt"
	"unsafe"

	"google.golang.org/protobuf/proto"

	dpfpb "github.com/google/distributed_point_functions/dpf/distributed_point_function_go_proto"
	pb "github.com/google/privacy-sandbox-aggregation-service/pipeline/crypto_go_proto"
)

// DefaultElementBitSize is the default element size for generating the DPF keys.
const DefaultElementBitSize = 64

func freeCBytes(cb C.struct_CBytes) {
	C.free(unsafe.Pointer(cb.c))
}

func createCParams(params []*dpfpb.DpfParameters) (*C.struct_CBytes, []unsafe.Pointer, error) {
	paramsLen := len(params)
	cParamPointers := make([]unsafe.Pointer, paramsLen)
	cParams := (*C.struct_CBytes)(C.malloc(C.sizeof_struct_CBytes * C.uint64_t(paramsLen)))
	pSlice := (*[1 << 30]C.struct_CBytes)(unsafe.Pointer(cParams))[:paramsLen:paramsLen]
	for i, p := range params {
		bParam, err := proto.Marshal(p)
		if err != nil {
			return nil, nil, err
		}
		cParamPointers[i] = C.CBytes(bParam)
		pSlice[i] = C.struct_CBytes{c: (*C.char)(cParamPointers[i]), l: C.int(len(bParam))}
	}
	return cParams, cParamPointers, nil
}

func freeCParams(cParams *C.struct_CBytes, cParamPointers []unsafe.Pointer) {
	for _, p := range cParamPointers {
		C.free(p)
	}
	C.free(unsafe.Pointer(cParams))
}

// GenerateKeys generates a pair of DpfKeys for given parameters.
func GenerateKeys(params []*dpfpb.DpfParameters, alpha uint64, betas []uint64) (*dpfpb.DpfKey, *dpfpb.DpfKey, error) {
	cParamsSize := C.int64_t(len(params))
	cParams, cParamPointers, err := createCParams(params)
	defer freeCParams(cParams, cParamPointers)
	if err != nil {
		return nil, nil, err
	}

	betasSize := len(betas)
	cBetasSize := C.int64_t(betasSize)
	var betasPointer *uint64
	if betasSize > 0 {
		betasPointer = &betas[0]
	}

	cKey1 := C.struct_CBytes{}
	cKey2 := C.struct_CBytes{}
	errStr := C.struct_CBytes{}
	status := C.CGenerateKeys(cParams, cParamsSize, C.uint64_t(alpha), (*C.uint64_t)(unsafe.Pointer(betasPointer)), cBetasSize, &cKey1, &cKey2, &errStr)
	defer freeCBytes(cKey1)
	defer freeCBytes(cKey2)
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	key1 := &dpfpb.DpfKey{}
	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cKey1.c), cKey1.l), key1); err != nil {
		return nil, nil, err
	}

	key2 := &dpfpb.DpfKey{}
	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cKey2.c), cKey2.l), key2); err != nil {
		return nil, nil, err
	}

	return key1, key2, nil
}

// CreateEvaluationContext creates the context for expanding the vectors.
func CreateEvaluationContext(params []*dpfpb.DpfParameters, key *dpfpb.DpfKey) (*dpfpb.EvaluationContext, error) {
	cParamsSize := C.int64_t(len(params))
	cParams, cParamPointers, err := createCParams(params)
	defer freeCParams(cParams, cParamPointers)
	if err != nil {
		return nil, err
	}

	bKey, err := proto.Marshal(key)
	if err != nil {
		return nil, err
	}
	cKey := C.struct_CBytes{c: (*C.char)(C.CBytes(bKey)), l: C.int(len(bKey))}
	defer freeCBytes(cKey)

	cEvalCtx := C.struct_CBytes{}
	errStr := C.struct_CBytes{}
	status := C.CCreateEvaluationContext(cParams, cParamsSize, &cKey, &cEvalCtx, &errStr)
	defer freeCBytes(cEvalCtx)
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	evalCtx := &dpfpb.EvaluationContext{}
	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cEvalCtx.c), cEvalCtx.l), evalCtx); err != nil {
		return nil, err
	}
	return evalCtx, nil
}

// EvaluateNext64 evaluates the given DPF key in the evaluation context with the specified configuration.
func EvaluateNext64(prefixes []uint64, evalCtx *dpfpb.EvaluationContext) ([]uint64, error) {
	pSize := len(prefixes)
	cPrefixesSize := C.int64_t(pSize)
	var prefixesPointer *uint64
	if pSize > 0 {
		prefixesPointer = &prefixes[0]
	}

	bEvalCtx, err := proto.Marshal(evalCtx)
	if err != nil {
		return nil, err
	}
	cEvalCtx := C.struct_CBytes{c: (*C.char)(C.CBytes(bEvalCtx)), l: C.int(len(bEvalCtx))}
	outExpanded := C.struct_CUInt64Vec{}
	errStr := C.struct_CBytes{}
	status := C.CEvaluateNext64((*C.uint64_t)(unsafe.Pointer(prefixesPointer)), cPrefixesSize, &cEvalCtx, &outExpanded, &errStr)
	defer freeCBytes(cEvalCtx)
	defer C.free(unsafe.Pointer(outExpanded.vec))
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cEvalCtx.c), cEvalCtx.l), evalCtx); err != nil {
		return nil, err
	}

	const maxLen = 1 << 30
	vecLen := uint64(outExpanded.vec_size)
	if vecLen > maxLen {
		return nil, fmt.Errorf("vector length %d should not exceed %d", vecLen, maxLen)
	}
	es := (*[maxLen]C.uint64_t)(unsafe.Pointer(outExpanded.vec))[:vecLen:vecLen]
	expanded := make([]uint64, vecLen)
	for i := uint64(0); i < uint64(vecLen); i++ {
		expanded[i] = uint64(es[i])
	}
	return expanded, nil
}

// EvaluateUntil64 evaluates the given DPF key in the evaluation context to a certain level of hierarchy.
func EvaluateUntil64(hierarchyLevel int, prefixes []uint64, evalCtx *dpfpb.EvaluationContext) ([]uint64, error) {
	pSize := len(prefixes)
	cPrefixesSize := C.int64_t(pSize)
	var prefixesPointer *uint64
	if pSize > 0 {
		prefixesPointer = &prefixes[0]
	}

	bEvalCtx, err := proto.Marshal(evalCtx)
	if err != nil {
		return nil, err
	}
	cEvalCtx := C.struct_CBytes{c: (*C.char)(C.CBytes(bEvalCtx)), l: C.int(len(bEvalCtx))}
	outExpanded := C.struct_CUInt64Vec{}
	errStr := C.struct_CBytes{}
	status := C.CEvaluateUntil64(C.int(hierarchyLevel), (*C.uint64_t)(unsafe.Pointer(prefixesPointer)), cPrefixesSize, &cEvalCtx, &outExpanded, &errStr)
	defer freeCBytes(cEvalCtx)
	defer C.free(unsafe.Pointer(outExpanded.vec))
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cEvalCtx.c), cEvalCtx.l), evalCtx); err != nil {
		return nil, err
	}

	const maxLen = 1 << 30
	vecLen := uint64(outExpanded.vec_size)
	if vecLen > maxLen {
		return nil, fmt.Errorf("vector length %d should not exceed %d", vecLen, maxLen)
	}
	es := (*[maxLen]C.uint64_t)(unsafe.Pointer(outExpanded.vec))[:vecLen:vecLen]
	expanded := make([]uint64, vecLen)
	for i := uint64(0); i < uint64(vecLen); i++ {
		expanded[i] = uint64(es[i])
	}
	return expanded, nil
}

// CalculateBucketID gets the bucket ID for values in the expanded vectors for certain level of hierarchy.
// If previousLevel = 1, the DPF key has not been evaluated yet:
// http://github.com/google/distributed_point_functions/dpf/distributed_point_function.cc?l=730&rcl=396584858
func CalculateBucketID(params *pb.IncrementalDpfParameters, prefixes *pb.HierarchicalPrefixes, levels []int32, previousLevel int32) ([]uint64, error) {
	if err := CheckExpansionParameters(params, prefixes, levels, previousLevel); err != nil {
		return nil, err
	}

	// For direct expansion, return empty slice to avoid generating extra data.
	// Because in this case, the bucket ID equals the vector index.
	if len(levels) == 1 && previousLevel == -1 {
		return nil, nil
	}

	var prefixBitSize int32
	if len(levels) > 1 {
		prefixBitSize = params.Params[levels[len(levels)-2]].GetLogDomainSize()
	} else {
		prefixBitSize = params.Params[previousLevel].GetLogDomainSize()
	}
	finalBitSize := params.Params[levels[len(levels)-1]].GetLogDomainSize()
	finalPrefixes := prefixes.Prefixes[len(prefixes.Prefixes)-1]

	expansionBits := finalBitSize - prefixBitSize
	expansionSize := uint64(1) << expansionBits
	ids := make([]uint64, uint64(len(finalPrefixes.Prefix))*expansionSize)
	i := uint64(0)
	for _, p := range finalPrefixes.Prefix {
		prefix := p << uint64(expansionBits)
		for j := uint64(0); j < expansionSize; j++ {
			ids[i] = prefix | j
			i++
		}
	}
	return ids, nil
}

func validateLevels(levels []int32) error {
	if len(levels) == 0 {
		return errors.New("expect non-empty levels")
	}
	level := levels[0]
	if level < 0 {
		return fmt.Errorf("expect non-negative level, got %d", level)
	}
	for i := 1; i < len(levels); i++ {
		if levels[i] <= level {
			return fmt.Errorf("expect levels in ascending order, got %v", levels)
		}
		level = levels[i]
	}
	return nil
}

// CheckExpansionParameters checks if the DPF parameters and prefixes are valid for the hierarchical expansion.
func CheckExpansionParameters(params *pb.IncrementalDpfParameters, prefixes *pb.HierarchicalPrefixes, levels []int32, previousLevel int32) error {
	paramsLen := len(params.Params)
	if paramsLen == 0 {
		return errors.New("empty dpf parameters")
	}

	if err := validateLevels(levels); err != nil {
		return err
	}
	if previousLevel < -1 {
		return fmt.Errorf("expect previousLevel >= -1, got %d", previousLevel)
	}
	if previousLevel >= levels[0] {
		return fmt.Errorf("expect previousLevel < levels[0] = %d, got %d", levels[0], previousLevel)
	}

	prefixesLen := len(prefixes.Prefixes)
	if prefixesLen != len(levels) {
		return fmt.Errorf("expansion level size should equal prefixes size %d, got %d", prefixesLen, len(levels))
	}
	if paramsLen < prefixesLen {
		return fmt.Errorf("expect prefixes size <= DPF parameter size %d, got %d", paramsLen, prefixesLen)
	}

	for i, p := range prefixes.Prefixes {
		if i == 0 && previousLevel < 0 && len(p.Prefix) != 0 {
			return fmt.Errorf("prefixes should be empty for the first level expansion, got %s", prefixes.String())
		}

		if !(i == 0 && previousLevel < 0) && len(p.Prefix) == 0 {
			return fmt.Errorf("prefix cannot be empty except for the first level expansion, got %s with previous level %d", prefixes.String(), previousLevel)
		}
	}
	return nil
}

// CheckOldExpansionParameters checks if the DPF parameters and prefixes defined in the old pipeline are valid for the hierarchical expansion.
func CheckOldExpansionParameters(params *pb.IncrementalDpfParameters, prefixes *pb.HierarchicalPrefixes) error {
	paramsLen := len(params.Params)
	if paramsLen == 0 {
		return errors.New("empty dpf parameters")
	}

	prefixesLen := len(prefixes.Prefixes)
	if paramsLen != prefixesLen {
		return fmt.Errorf("dpf parameter size should equal prefixes size %d, got %d", prefixesLen, paramsLen)
	}

	if len(prefixes.Prefixes[0].Prefix) != 0 {
		return fmt.Errorf("prefixes should be empty for the first level expansion, got %s", prefixes.Prefixes[0].String())
	}

	for i, p := range prefixes.Prefixes {
		if i > 0 && len(p.Prefix) == 0 {
			return fmt.Errorf("prefix cannot be empty except for the top level expansion")
		}
	}
	return nil
}
<commit_msg>Remove an unused function for compatibility of the old version.<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 incrementaldpf calls the C++ incremental DPF functions from Go.
package incrementaldpf

// #cgo CFLAGS: -g -Wall
// #include <stdbool.h>
// #include <stdlib.h>
// #include "pipeline/distributed_point_function_c_bridge.h"
import (
	"C"
)

import (
	"errors"
	"fmt"
	"unsafe"

	"google.golang.org/protobuf/proto"

	dpfpb "github.com/google/distributed_point_functions/dpf/distributed_point_function_go_proto"
	pb "github.com/google/privacy-sandbox-aggregation-service/pipeline/crypto_go_proto"
)

// DefaultElementBitSize is the default element size for generating the DPF keys.
const DefaultElementBitSize = 64

func freeCBytes(cb C.struct_CBytes) {
	C.free(unsafe.Pointer(cb.c))
}

func createCParams(params []*dpfpb.DpfParameters) (*C.struct_CBytes, []unsafe.Pointer, error) {
	paramsLen := len(params)
	cParamPointers := make([]unsafe.Pointer, paramsLen)
	cParams := (*C.struct_CBytes)(C.malloc(C.sizeof_struct_CBytes * C.uint64_t(paramsLen)))
	pSlice := (*[1 << 30]C.struct_CBytes)(unsafe.Pointer(cParams))[:paramsLen:paramsLen]
	for i, p := range params {
		bParam, err := proto.Marshal(p)
		if err != nil {
			return nil, nil, err
		}
		cParamPointers[i] = C.CBytes(bParam)
		pSlice[i] = C.struct_CBytes{c: (*C.char)(cParamPointers[i]), l: C.int(len(bParam))}
	}
	return cParams, cParamPointers, nil
}

func freeCParams(cParams *C.struct_CBytes, cParamPointers []unsafe.Pointer) {
	for _, p := range cParamPointers {
		C.free(p)
	}
	C.free(unsafe.Pointer(cParams))
}

// GenerateKeys generates a pair of DpfKeys for given parameters.
func GenerateKeys(params []*dpfpb.DpfParameters, alpha uint64, betas []uint64) (*dpfpb.DpfKey, *dpfpb.DpfKey, error) {
	cParamsSize := C.int64_t(len(params))
	cParams, cParamPointers, err := createCParams(params)
	defer freeCParams(cParams, cParamPointers)
	if err != nil {
		return nil, nil, err
	}

	betasSize := len(betas)
	cBetasSize := C.int64_t(betasSize)
	var betasPointer *uint64
	if betasSize > 0 {
		betasPointer = &betas[0]
	}

	cKey1 := C.struct_CBytes{}
	cKey2 := C.struct_CBytes{}
	errStr := C.struct_CBytes{}
	status := C.CGenerateKeys(cParams, cParamsSize, C.uint64_t(alpha), (*C.uint64_t)(unsafe.Pointer(betasPointer)), cBetasSize, &cKey1, &cKey2, &errStr)
	defer freeCBytes(cKey1)
	defer freeCBytes(cKey2)
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	key1 := &dpfpb.DpfKey{}
	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cKey1.c), cKey1.l), key1); err != nil {
		return nil, nil, err
	}

	key2 := &dpfpb.DpfKey{}
	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cKey2.c), cKey2.l), key2); err != nil {
		return nil, nil, err
	}

	return key1, key2, nil
}

// CreateEvaluationContext creates the context for expanding the vectors.
func CreateEvaluationContext(params []*dpfpb.DpfParameters, key *dpfpb.DpfKey) (*dpfpb.EvaluationContext, error) {
	cParamsSize := C.int64_t(len(params))
	cParams, cParamPointers, err := createCParams(params)
	defer freeCParams(cParams, cParamPointers)
	if err != nil {
		return nil, err
	}

	bKey, err := proto.Marshal(key)
	if err != nil {
		return nil, err
	}
	cKey := C.struct_CBytes{c: (*C.char)(C.CBytes(bKey)), l: C.int(len(bKey))}
	defer freeCBytes(cKey)

	cEvalCtx := C.struct_CBytes{}
	errStr := C.struct_CBytes{}
	status := C.CCreateEvaluationContext(cParams, cParamsSize, &cKey, &cEvalCtx, &errStr)
	defer freeCBytes(cEvalCtx)
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	evalCtx := &dpfpb.EvaluationContext{}
	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cEvalCtx.c), cEvalCtx.l), evalCtx); err != nil {
		return nil, err
	}
	return evalCtx, nil
}

// EvaluateNext64 evaluates the given DPF key in the evaluation context with the specified configuration.
func EvaluateNext64(prefixes []uint64, evalCtx *dpfpb.EvaluationContext) ([]uint64, error) {
	pSize := len(prefixes)
	cPrefixesSize := C.int64_t(pSize)
	var prefixesPointer *uint64
	if pSize > 0 {
		prefixesPointer = &prefixes[0]
	}

	bEvalCtx, err := proto.Marshal(evalCtx)
	if err != nil {
		return nil, err
	}
	cEvalCtx := C.struct_CBytes{c: (*C.char)(C.CBytes(bEvalCtx)), l: C.int(len(bEvalCtx))}
	outExpanded := C.struct_CUInt64Vec{}
	errStr := C.struct_CBytes{}
	status := C.CEvaluateNext64((*C.uint64_t)(unsafe.Pointer(prefixesPointer)), cPrefixesSize, &cEvalCtx, &outExpanded, &errStr)
	defer freeCBytes(cEvalCtx)
	defer C.free(unsafe.Pointer(outExpanded.vec))
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cEvalCtx.c), cEvalCtx.l), evalCtx); err != nil {
		return nil, err
	}

	const maxLen = 1 << 30
	vecLen := uint64(outExpanded.vec_size)
	if vecLen > maxLen {
		return nil, fmt.Errorf("vector length %d should not exceed %d", vecLen, maxLen)
	}
	es := (*[maxLen]C.uint64_t)(unsafe.Pointer(outExpanded.vec))[:vecLen:vecLen]
	expanded := make([]uint64, vecLen)
	for i := uint64(0); i < uint64(vecLen); i++ {
		expanded[i] = uint64(es[i])
	}
	return expanded, nil
}

// EvaluateUntil64 evaluates the given DPF key in the evaluation context to a certain level of hierarchy.
func EvaluateUntil64(hierarchyLevel int, prefixes []uint64, evalCtx *dpfpb.EvaluationContext) ([]uint64, error) {
	pSize := len(prefixes)
	cPrefixesSize := C.int64_t(pSize)
	var prefixesPointer *uint64
	if pSize > 0 {
		prefixesPointer = &prefixes[0]
	}

	bEvalCtx, err := proto.Marshal(evalCtx)
	if err != nil {
		return nil, err
	}
	cEvalCtx := C.struct_CBytes{c: (*C.char)(C.CBytes(bEvalCtx)), l: C.int(len(bEvalCtx))}
	outExpanded := C.struct_CUInt64Vec{}
	errStr := C.struct_CBytes{}
	status := C.CEvaluateUntil64(C.int(hierarchyLevel), (*C.uint64_t)(unsafe.Pointer(prefixesPointer)), cPrefixesSize, &cEvalCtx, &outExpanded, &errStr)
	defer freeCBytes(cEvalCtx)
	defer C.free(unsafe.Pointer(outExpanded.vec))
	defer freeCBytes(errStr)
	if status != 0 {
		return nil, errors.New(C.GoStringN(errStr.c, errStr.l))
	}

	if err := proto.Unmarshal(C.GoBytes(unsafe.Pointer(cEvalCtx.c), cEvalCtx.l), evalCtx); err != nil {
		return nil, err
	}

	const maxLen = 1 << 30
	vecLen := uint64(outExpanded.vec_size)
	if vecLen > maxLen {
		return nil, fmt.Errorf("vector length %d should not exceed %d", vecLen, maxLen)
	}
	es := (*[maxLen]C.uint64_t)(unsafe.Pointer(outExpanded.vec))[:vecLen:vecLen]
	expanded := make([]uint64, vecLen)
	for i := uint64(0); i < uint64(vecLen); i++ {
		expanded[i] = uint64(es[i])
	}
	return expanded, nil
}

// CalculateBucketID gets the bucket ID for values in the expanded vectors for certain level of hierarchy.
// If previousLevel = 1, the DPF key has not been evaluated yet:
// http://github.com/google/distributed_point_functions/dpf/distributed_point_function.cc?l=730&rcl=396584858
func CalculateBucketID(params *pb.IncrementalDpfParameters, prefixes *pb.HierarchicalPrefixes, levels []int32, previousLevel int32) ([]uint64, error) {
	if err := CheckExpansionParameters(params, prefixes, levels, previousLevel); err != nil {
		return nil, err
	}

	// For direct expansion, return empty slice to avoid generating extra data.
	// Because in this case, the bucket ID equals the vector index.
	if len(levels) == 1 && previousLevel == -1 {
		return nil, nil
	}

	var prefixBitSize int32
	if len(levels) > 1 {
		prefixBitSize = params.Params[levels[len(levels)-2]].GetLogDomainSize()
	} else {
		prefixBitSize = params.Params[previousLevel].GetLogDomainSize()
	}
	finalBitSize := params.Params[levels[len(levels)-1]].GetLogDomainSize()
	finalPrefixes := prefixes.Prefixes[len(prefixes.Prefixes)-1]

	expansionBits := finalBitSize - prefixBitSize
	expansionSize := uint64(1) << expansionBits
	ids := make([]uint64, uint64(len(finalPrefixes.Prefix))*expansionSize)
	i := uint64(0)
	for _, p := range finalPrefixes.Prefix {
		prefix := p << uint64(expansionBits)
		for j := uint64(0); j < expansionSize; j++ {
			ids[i] = prefix | j
			i++
		}
	}
	return ids, nil
}

func validateLevels(levels []int32) error {
	if len(levels) == 0 {
		return errors.New("expect non-empty levels")
	}
	level := levels[0]
	if level < 0 {
		return fmt.Errorf("expect non-negative level, got %d", level)
	}
	for i := 1; i < len(levels); i++ {
		if levels[i] <= level {
			return fmt.Errorf("expect levels in ascending order, got %v", levels)
		}
		level = levels[i]
	}
	return nil
}

// CheckExpansionParameters checks if the DPF parameters and prefixes are valid for the hierarchical expansion.
func CheckExpansionParameters(params *pb.IncrementalDpfParameters, prefixes *pb.HierarchicalPrefixes, levels []int32, previousLevel int32) error {
	paramsLen := len(params.Params)
	if paramsLen == 0 {
		return errors.New("empty dpf parameters")
	}

	if err := validateLevels(levels); err != nil {
		return err
	}
	if previousLevel < -1 {
		return fmt.Errorf("expect previousLevel >= -1, got %d", previousLevel)
	}
	if previousLevel >= levels[0] {
		return fmt.Errorf("expect previousLevel < levels[0] = %d, got %d", levels[0], previousLevel)
	}

	prefixesLen := len(prefixes.Prefixes)
	if prefixesLen != len(levels) {
		return fmt.Errorf("expansion level size should equal prefixes size %d, got %d", prefixesLen, len(levels))
	}
	if paramsLen < prefixesLen {
		return fmt.Errorf("expect prefixes size <= DPF parameter size %d, got %d", paramsLen, prefixesLen)
	}

	for i, p := range prefixes.Prefixes {
		if i == 0 && previousLevel < 0 && len(p.Prefix) != 0 {
			return fmt.Errorf("prefixes should be empty for the first level expansion, got %s", prefixes.String())
		}

		if !(i == 0 && previousLevel < 0) && len(p.Prefix) == 0 {
			return fmt.Errorf("prefix cannot be empty except for the first level expansion, got %s with previous level %d", prefixes.String(), previousLevel)
		}
	}
	return nil
}
<|endoftext|>


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

package optimize

import (
	"math"

	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

const maxNewtonModifications = 20

// Newton implements a modified Newton's method for Hessian-based unconstrained
// minimization. It applies regularization when the Hessian is not positive
// definite, and it can converge to a local minimum from any starting point.
//
// Newton iteratively forms a quadratic model to the objective function f and
// tries to minimize this approximate model. It generates a sequence of
// locations x_k by means of
//  solve H_k d_k = -∇f_k for d_k,
//  x_{k+1} = x_k + α_k d_k,
// where H_k is the Hessian matrix of f at x_k and α_k is a step size found by
// a line search.
//
// Away from a minimizer H_k may not be positive definite and d_k may not be a
// descent direction. Newton implements a Hessian modification strategy that
// adds successively larger multiples of identity to H_k until it becomes
// positive definite. Note that the repeated trial factorization of the
// modified Hessian involved in this process can be computationally expensive.
//
// If the Hessian matrix cannot be formed explicitly or if the computational
// cost of its factorization is prohibitive, BFGS or L-BFGS quasi-Newton method
// can be used instead.
type Newton struct {
	// Linesearcher is used for selecting suitable steps along the descent
	// direction d. Accepted steps should satisfy at least one of the Wolfe,
	// Goldstein or Armijo conditions.
	// If Linesearcher == nil, an appropriate default is chosen.
	Linesearcher Linesearcher
	// Increase is the factor by which a scalar tau is successively increased
	// so that (H + tau*I) is positive definite. Larger values reduce the
	// number of trial Hessian factorizations, but also reduce the second-order
	// information in H.
	// Increase must be greater than 1. If Increase is 0, it is defaulted to 5.
	Increase float64

	ls *LinesearchMethod

	hess *mat64.SymDense // Storage for a copy of the Hessian matrix.
	chol *mat64.Cholesky // Storage for the Cholesky factorization.
	tau  float64
}

func (n *Newton) Init(loc *Location, xNext []float64) (EvaluationType, IterationType, error) {
	if n.Increase == 0 {
		n.Increase = 5
	}
	if n.Increase <= 1 {
		panic("optimize: Newton.Increase must be greater than 1")
	}
	if n.Linesearcher == nil {
		n.Linesearcher = &Bisection{}
	}
	if n.ls == nil {
		n.ls = &LinesearchMethod{}
	}
	n.ls.Linesearcher = n.Linesearcher
	n.ls.NextDirectioner = n

	return n.ls.Init(loc, xNext)
}

func (n *Newton) Iterate(loc *Location, xNext []float64) (EvaluationType, IterationType, error) {
	return n.ls.Iterate(loc, xNext)
}

func (n *Newton) InitDirection(loc *Location, dir []float64) (stepSize float64) {
	dim := len(loc.X)
	if n.chol == nil {
		n.chol = &mat64.Cholesky{}
	}
	n.hess = resizeSymDense(n.hess, dim)
	n.tau = 0
	return n.NextDirection(loc, dir)
}

func (n *Newton) NextDirection(loc *Location, dir []float64) (stepSize float64) {
	// This method implements Algorithm 3.3 (Cholesky with Added Multiple of
	// the Identity) from Nocedal, Wright (2006), 2nd edition.

	dim := len(loc.X)
	n.hess.CopySym(loc.Hessian)

	// Find the smallest diagonal entry of the Hesssian.
	minA := n.hess.At(0, 0)
	for i := 1; i < dim; i++ {
		a := n.hess.At(i, i)
		if a < minA {
			minA = a
		}
	}
	// If the smallest diagonal entry is positive, the Hessian may be positive
	// definite, and so first attempt to apply the Cholesky factorization to
	// the un-modified Hessian. If the smallest entry is negative, use the
	// final tau from the last iteration if regularization was needed,
	// otherwise guess an appropriate value for tau.
	if minA > 0 {
		n.tau = 0
	} else if n.tau == 0 {
		n.tau = -minA + 0.001
	}

	for k := 0; k < maxNewtonModifications; k++ {
		if n.tau != 0 {
			// Add a multiple of identity to the Hessian.
			for i := 0; i < dim; i++ {
				n.hess.SetSym(i, i, loc.Hessian.At(i, i)+n.tau)
			}
		}
		// Try to apply the Cholesky factorization.
		pd := n.chol.Factorize(n.hess)
		if pd {
			d := mat64.NewVector(dim, dir)
			// Store the solution in d's backing array, dir.
			d.SolveCholeskyVec(n.chol, mat64.NewVector(dim, loc.Gradient))
			floats.Scale(-1, dir)
			return 1
		}
		// Modified Hessian is not PD, so increase tau.
		n.tau = math.Max(n.Increase*n.tau, 0.001)
	}

	// Hessian modification failed to get a PD matrix. Return the negative
	// gradient as the descent direction.
	copy(dir, loc.Gradient)
	floats.Scale(-1, dir)
	return 1
}

func (n *Newton) Needs() struct {
	Gradient bool
	Hessian  bool
} {
	return struct {
		Gradient bool
		Hessian  bool
	}{true, true}
}
<commit_msg>optimize: make chol field concrete type<commit_after>// Copyright ©2015 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package optimize

import (
	"math"

	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

const maxNewtonModifications = 20

// Newton implements a modified Newton's method for Hessian-based unconstrained
// minimization. It applies regularization when the Hessian is not positive
// definite, and it can converge to a local minimum from any starting point.
//
// Newton iteratively forms a quadratic model to the objective function f and
// tries to minimize this approximate model. It generates a sequence of
// locations x_k by means of
//  solve H_k d_k = -∇f_k for d_k,
//  x_{k+1} = x_k + α_k d_k,
// where H_k is the Hessian matrix of f at x_k and α_k is a step size found by
// a line search.
//
// Away from a minimizer H_k may not be positive definite and d_k may not be a
// descent direction. Newton implements a Hessian modification strategy that
// adds successively larger multiples of identity to H_k until it becomes
// positive definite. Note that the repeated trial factorization of the
// modified Hessian involved in this process can be computationally expensive.
//
// If the Hessian matrix cannot be formed explicitly or if the computational
// cost of its factorization is prohibitive, BFGS or L-BFGS quasi-Newton method
// can be used instead.
type Newton struct {
	// Linesearcher is used for selecting suitable steps along the descent
	// direction d. Accepted steps should satisfy at least one of the Wolfe,
	// Goldstein or Armijo conditions.
	// If Linesearcher == nil, an appropriate default is chosen.
	Linesearcher Linesearcher
	// Increase is the factor by which a scalar tau is successively increased
	// so that (H + tau*I) is positive definite. Larger values reduce the
	// number of trial Hessian factorizations, but also reduce the second-order
	// information in H.
	// Increase must be greater than 1. If Increase is 0, it is defaulted to 5.
	Increase float64

	ls *LinesearchMethod

	hess *mat64.SymDense // Storage for a copy of the Hessian matrix.
	chol mat64.Cholesky  // Storage for the Cholesky factorization.
	tau  float64
}

func (n *Newton) Init(loc *Location, xNext []float64) (EvaluationType, IterationType, error) {
	if n.Increase == 0 {
		n.Increase = 5
	}
	if n.Increase <= 1 {
		panic("optimize: Newton.Increase must be greater than 1")
	}
	if n.Linesearcher == nil {
		n.Linesearcher = &Bisection{}
	}
	if n.ls == nil {
		n.ls = &LinesearchMethod{}
	}
	n.ls.Linesearcher = n.Linesearcher
	n.ls.NextDirectioner = n

	return n.ls.Init(loc, xNext)
}

func (n *Newton) Iterate(loc *Location, xNext []float64) (EvaluationType, IterationType, error) {
	return n.ls.Iterate(loc, xNext)
}

func (n *Newton) InitDirection(loc *Location, dir []float64) (stepSize float64) {
	dim := len(loc.X)
	n.hess = resizeSymDense(n.hess, dim)
	n.tau = 0
	return n.NextDirection(loc, dir)
}

func (n *Newton) NextDirection(loc *Location, dir []float64) (stepSize float64) {
	// This method implements Algorithm 3.3 (Cholesky with Added Multiple of
	// the Identity) from Nocedal, Wright (2006), 2nd edition.

	dim := len(loc.X)
	n.hess.CopySym(loc.Hessian)

	// Find the smallest diagonal entry of the Hesssian.
	minA := n.hess.At(0, 0)
	for i := 1; i < dim; i++ {
		a := n.hess.At(i, i)
		if a < minA {
			minA = a
		}
	}
	// If the smallest diagonal entry is positive, the Hessian may be positive
	// definite, and so first attempt to apply the Cholesky factorization to
	// the un-modified Hessian. If the smallest entry is negative, use the
	// final tau from the last iteration if regularization was needed,
	// otherwise guess an appropriate value for tau.
	if minA > 0 {
		n.tau = 0
	} else if n.tau == 0 {
		n.tau = -minA + 0.001
	}

	for k := 0; k < maxNewtonModifications; k++ {
		if n.tau != 0 {
			// Add a multiple of identity to the Hessian.
			for i := 0; i < dim; i++ {
				n.hess.SetSym(i, i, loc.Hessian.At(i, i)+n.tau)
			}
		}
		// Try to apply the Cholesky factorization.
		pd := n.chol.Factorize(n.hess)
		if pd {
			d := mat64.NewVector(dim, dir)
			// Store the solution in d's backing array, dir.
			d.SolveCholeskyVec(&n.chol, mat64.NewVector(dim, loc.Gradient))
			floats.Scale(-1, dir)
			return 1
		}
		// Modified Hessian is not PD, so increase tau.
		n.tau = math.Max(n.Increase*n.tau, 0.001)
	}

	// Hessian modification failed to get a PD matrix. Return the negative
	// gradient as the descent direction.
	copy(dir, loc.Gradient)
	floats.Scale(-1, dir)
	return 1
}

func (n *Newton) Needs() struct {
	Gradient bool
	Hessian  bool
} {
	return struct {
		Gradient bool
		Hessian  bool
	}{true, true}
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 action

import (
	"fmt"
	"testing"

	"helm.sh/helm/v3/pkg/chart"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	kubefake "helm.sh/helm/v3/pkg/kube/fake"
	"helm.sh/helm/v3/pkg/release"
	"helm.sh/helm/v3/pkg/time"
)

func upgradeAction(t *testing.T) *Upgrade {
	config := actionConfigFixture(t)
	upAction := NewUpgrade(config)
	upAction.Namespace = "spaced"

	return upAction
}

func TestUpgradeRelease_Wait(t *testing.T) {
	is := assert.New(t)
	req := require.New(t)

	upAction := upgradeAction(t)
	rel := releaseStub()
	rel.Name = "come-fail-away"
	rel.Info.Status = release.StatusDeployed
	upAction.cfg.Releases.Create(rel)

	failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
	failer.WaitError = fmt.Errorf("I timed out")
	upAction.cfg.KubeClient = failer
	upAction.Wait = true
	vals := map[string]interface{}{}

	res, err := upAction.Run(rel.Name, buildChart(), vals)
	req.Error(err)
	is.Contains(res.Info.Description, "I timed out")
	is.Equal(res.Info.Status, release.StatusFailed)
}

func TestUpgradeRelease_CleanupOnFail(t *testing.T) {
	is := assert.New(t)
	req := require.New(t)

	upAction := upgradeAction(t)
	rel := releaseStub()
	rel.Name = "come-fail-away"
	rel.Info.Status = release.StatusDeployed
	upAction.cfg.Releases.Create(rel)

	failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
	failer.WaitError = fmt.Errorf("I timed out")
	failer.DeleteError = fmt.Errorf("I tried to delete nil")
	upAction.cfg.KubeClient = failer
	upAction.Wait = true
	upAction.CleanupOnFail = true
	vals := map[string]interface{}{}

	res, err := upAction.Run(rel.Name, buildChart(), vals)
	req.Error(err)
	is.NotContains(err.Error(), "unable to cleanup resources")
	is.Contains(res.Info.Description, "I timed out")
	is.Equal(res.Info.Status, release.StatusFailed)
}

func TestUpgradeRelease_Atomic(t *testing.T) {
	is := assert.New(t)
	req := require.New(t)

	t.Run("atomic rollback succeeds", func(t *testing.T) {
		upAction := upgradeAction(t)

		rel := releaseStub()
		rel.Name = "nuketown"
		rel.Info.Status = release.StatusDeployed
		upAction.cfg.Releases.Create(rel)

		failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
		// We can't make Update error because then the rollback won't work
		failer.WatchUntilReadyError = fmt.Errorf("arming key removed")
		upAction.cfg.KubeClient = failer
		upAction.Atomic = true
		vals := map[string]interface{}{}

		res, err := upAction.Run(rel.Name, buildChart(), vals)
		req.Error(err)
		is.Contains(err.Error(), "arming key removed")
		is.Contains(err.Error(), "atomic")

		// Now make sure it is actually upgraded
		updatedRes, err := upAction.cfg.Releases.Get(res.Name, 3)
		is.NoError(err)
		// Should have rolled back to the previous
		is.Equal(updatedRes.Info.Status, release.StatusDeployed)
	})

	t.Run("atomic uninstall fails", func(t *testing.T) {
		upAction := upgradeAction(t)
		rel := releaseStub()
		rel.Name = "fallout"
		rel.Info.Status = release.StatusDeployed
		upAction.cfg.Releases.Create(rel)

		failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
		failer.UpdateError = fmt.Errorf("update fail")
		upAction.cfg.KubeClient = failer
		upAction.Atomic = true
		vals := map[string]interface{}{}

		_, err := upAction.Run(rel.Name, buildChart(), vals)
		req.Error(err)
		is.Contains(err.Error(), "update fail")
		is.Contains(err.Error(), "an error occurred while rolling back the release")
	})
}

func TestUpgradeRelease_ReuseValues(t *testing.T) {
	is := assert.New(t)

	t.Run("reuse values should work with values", func(t *testing.T) {
		upAction := upgradeAction(t)

		existingValues := map[string]interface{}{
			"name":        "value",
			"maxHeapSize": "128m",
			"replicas":    2,
		}
		newValues := map[string]interface{}{
			"name":        "newValue",
			"maxHeapSize": "512m",
			"cpu":         "12m",
		}
		expectedValues := map[string]interface{}{
			"name":        "newValue",
			"maxHeapSize": "512m",
			"cpu":         "12m",
			"replicas":    2,
		}

		rel := releaseStub()
		rel.Name = "nuketown"
		rel.Info.Status = release.StatusDeployed
		rel.Config = existingValues

		err := upAction.cfg.Releases.Create(rel)
		is.NoError(err)

		upAction.ReuseValues = true
		// setting newValues and upgrading
		res, err := upAction.Run(rel.Name, buildChart(), newValues)
		is.NoError(err)

		// Now make sure it is actually upgraded
		updatedRes, err := upAction.cfg.Releases.Get(res.Name, 2)
		is.NoError(err)

		if updatedRes == nil {
			is.Fail("Updated Release is nil")
			return
		}
		is.Equal(release.StatusDeployed, updatedRes.Info.Status)
		is.Equal(expectedValues, updatedRes.Config)
	})

	t.Run("reuse values should not install disabled charts", func(t *testing.T) {
		upAction := upgradeAction(t)
		chartDefaultValues := map[string]interface{}{
			"subchart": map[string]interface{}{
				"enabled": true,
			},
		}
		dependency := chart.Dependency{
			Name:       "subchart",
			Version:    "0.1.0",
			Repository: "http://some-repo.com",
			Condition:  "subchart.enabled",
		}
		sampleChart := buildChart(
			withName("sample"),
			withValues(chartDefaultValues),
			withMetadataDependency(dependency),
		)
		now := time.Now()
		existingValues := map[string]interface{}{
			"subchart": map[string]interface{}{
				"enabled": false,
			},
		}
		rel := &release.Release{
			Name: "nuketown",
			Info: &release.Info{
				FirstDeployed: now,
				LastDeployed:  now,
				Status:        release.StatusDeployed,
				Description:   "Named Release Stub",
			},
			Chart:   sampleChart,
			Config:  existingValues,
			Version: 1,
		}
		err := upAction.cfg.Releases.Create(rel)
		is.NoError(err)

		upAction.ReuseValues = true
		sampleChartWithSubChart := buildChart(
			withName(sampleChart.Name()),
			withValues(sampleChart.Values),
			withDependency(withName("subchart")),
			withMetadataDependency(dependency),
		)
		// reusing values and upgrading
		res, err := upAction.Run(rel.Name, sampleChartWithSubChart, map[string]interface{}{})
		is.NoError(err)

		// Now get the upgraded release
		updatedRes, err := upAction.cfg.Releases.Get(res.Name, 2)
		is.NoError(err)

		if updatedRes == nil {
			is.Fail("Updated Release is nil")
			return
		}
		is.Equal(release.StatusDeployed, updatedRes.Info.Status)
		is.Equal(0, len(updatedRes.Chart.Dependencies()), "expected 0 dependencies")

		expectedValues := map[string]interface{}{
			"subchart": map[string]interface{}{
				"enabled": false,
			},
		}
		is.Equal(expectedValues, updatedRes.Config)
	})
}
<commit_msg>fix(helm): Added test for concurrent upgrades<commit_after>/*
Copyright The Helm 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 action

import (
	"fmt"
	"testing"

	"helm.sh/helm/v3/pkg/chart"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	kubefake "helm.sh/helm/v3/pkg/kube/fake"
	"helm.sh/helm/v3/pkg/release"
	"helm.sh/helm/v3/pkg/time"
)

func upgradeAction(t *testing.T) *Upgrade {
	config := actionConfigFixture(t)
	upAction := NewUpgrade(config)
	upAction.Namespace = "spaced"

	return upAction
}

func TestUpgradeRelease_Wait(t *testing.T) {
	is := assert.New(t)
	req := require.New(t)

	upAction := upgradeAction(t)
	rel := releaseStub()
	rel.Name = "come-fail-away"
	rel.Info.Status = release.StatusDeployed
	upAction.cfg.Releases.Create(rel)

	failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
	failer.WaitError = fmt.Errorf("I timed out")
	upAction.cfg.KubeClient = failer
	upAction.Wait = true
	vals := map[string]interface{}{}

	res, err := upAction.Run(rel.Name, buildChart(), vals)
	req.Error(err)
	is.Contains(res.Info.Description, "I timed out")
	is.Equal(res.Info.Status, release.StatusFailed)
}

func TestUpgradeRelease_CleanupOnFail(t *testing.T) {
	is := assert.New(t)
	req := require.New(t)

	upAction := upgradeAction(t)
	rel := releaseStub()
	rel.Name = "come-fail-away"
	rel.Info.Status = release.StatusDeployed
	upAction.cfg.Releases.Create(rel)

	failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
	failer.WaitError = fmt.Errorf("I timed out")
	failer.DeleteError = fmt.Errorf("I tried to delete nil")
	upAction.cfg.KubeClient = failer
	upAction.Wait = true
	upAction.CleanupOnFail = true
	vals := map[string]interface{}{}

	res, err := upAction.Run(rel.Name, buildChart(), vals)
	req.Error(err)
	is.NotContains(err.Error(), "unable to cleanup resources")
	is.Contains(res.Info.Description, "I timed out")
	is.Equal(res.Info.Status, release.StatusFailed)
}

func TestUpgradeRelease_Atomic(t *testing.T) {
	is := assert.New(t)
	req := require.New(t)

	t.Run("atomic rollback succeeds", func(t *testing.T) {
		upAction := upgradeAction(t)

		rel := releaseStub()
		rel.Name = "nuketown"
		rel.Info.Status = release.StatusDeployed
		upAction.cfg.Releases.Create(rel)

		failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
		// We can't make Update error because then the rollback won't work
		failer.WatchUntilReadyError = fmt.Errorf("arming key removed")
		upAction.cfg.KubeClient = failer
		upAction.Atomic = true
		vals := map[string]interface{}{}

		res, err := upAction.Run(rel.Name, buildChart(), vals)
		req.Error(err)
		is.Contains(err.Error(), "arming key removed")
		is.Contains(err.Error(), "atomic")

		// Now make sure it is actually upgraded
		updatedRes, err := upAction.cfg.Releases.Get(res.Name, 3)
		is.NoError(err)
		// Should have rolled back to the previous
		is.Equal(updatedRes.Info.Status, release.StatusDeployed)
	})

	t.Run("atomic uninstall fails", func(t *testing.T) {
		upAction := upgradeAction(t)
		rel := releaseStub()
		rel.Name = "fallout"
		rel.Info.Status = release.StatusDeployed
		upAction.cfg.Releases.Create(rel)

		failer := upAction.cfg.KubeClient.(*kubefake.FailingKubeClient)
		failer.UpdateError = fmt.Errorf("update fail")
		upAction.cfg.KubeClient = failer
		upAction.Atomic = true
		vals := map[string]interface{}{}

		_, err := upAction.Run(rel.Name, buildChart(), vals)
		req.Error(err)
		is.Contains(err.Error(), "update fail")
		is.Contains(err.Error(), "an error occurred while rolling back the release")
	})
}

func TestUpgradeRelease_ReuseValues(t *testing.T) {
	is := assert.New(t)

	t.Run("reuse values should work with values", func(t *testing.T) {
		upAction := upgradeAction(t)

		existingValues := map[string]interface{}{
			"name":        "value",
			"maxHeapSize": "128m",
			"replicas":    2,
		}
		newValues := map[string]interface{}{
			"name":        "newValue",
			"maxHeapSize": "512m",
			"cpu":         "12m",
		}
		expectedValues := map[string]interface{}{
			"name":        "newValue",
			"maxHeapSize": "512m",
			"cpu":         "12m",
			"replicas":    2,
		}

		rel := releaseStub()
		rel.Name = "nuketown"
		rel.Info.Status = release.StatusDeployed
		rel.Config = existingValues

		err := upAction.cfg.Releases.Create(rel)
		is.NoError(err)

		upAction.ReuseValues = true
		// setting newValues and upgrading
		res, err := upAction.Run(rel.Name, buildChart(), newValues)
		is.NoError(err)

		// Now make sure it is actually upgraded
		updatedRes, err := upAction.cfg.Releases.Get(res.Name, 2)
		is.NoError(err)

		if updatedRes == nil {
			is.Fail("Updated Release is nil")
			return
		}
		is.Equal(release.StatusDeployed, updatedRes.Info.Status)
		is.Equal(expectedValues, updatedRes.Config)
	})

	t.Run("reuse values should not install disabled charts", func(t *testing.T) {
		upAction := upgradeAction(t)
		chartDefaultValues := map[string]interface{}{
			"subchart": map[string]interface{}{
				"enabled": true,
			},
		}
		dependency := chart.Dependency{
			Name:       "subchart",
			Version:    "0.1.0",
			Repository: "http://some-repo.com",
			Condition:  "subchart.enabled",
		}
		sampleChart := buildChart(
			withName("sample"),
			withValues(chartDefaultValues),
			withMetadataDependency(dependency),
		)
		now := time.Now()
		existingValues := map[string]interface{}{
			"subchart": map[string]interface{}{
				"enabled": false,
			},
		}
		rel := &release.Release{
			Name: "nuketown",
			Info: &release.Info{
				FirstDeployed: now,
				LastDeployed:  now,
				Status:        release.StatusDeployed,
				Description:   "Named Release Stub",
			},
			Chart:   sampleChart,
			Config:  existingValues,
			Version: 1,
		}
		err := upAction.cfg.Releases.Create(rel)
		is.NoError(err)

		upAction.ReuseValues = true
		sampleChartWithSubChart := buildChart(
			withName(sampleChart.Name()),
			withValues(sampleChart.Values),
			withDependency(withName("subchart")),
			withMetadataDependency(dependency),
		)
		// reusing values and upgrading
		res, err := upAction.Run(rel.Name, sampleChartWithSubChart, map[string]interface{}{})
		is.NoError(err)

		// Now get the upgraded release
		updatedRes, err := upAction.cfg.Releases.Get(res.Name, 2)
		is.NoError(err)

		if updatedRes == nil {
			is.Fail("Updated Release is nil")
			return
		}
		is.Equal(release.StatusDeployed, updatedRes.Info.Status)
		is.Equal(0, len(updatedRes.Chart.Dependencies()), "expected 0 dependencies")

		expectedValues := map[string]interface{}{
			"subchart": map[string]interface{}{
				"enabled": false,
			},
		}
		is.Equal(expectedValues, updatedRes.Config)
	})
}

func TestUpgradeRelease_Pending(t *testing.T) {
	req := require.New(t)

	upAction := upgradeAction(t)
	rel := releaseStub()
	rel.Name = "come-fail-away"
	rel.Info.Status = release.StatusDeployed
	upAction.cfg.Releases.Create(rel)
	rel2 := releaseStub()
	rel2.Name = "come-fail-away"
	rel2.Info.Status = release.StatusPendingUpgrade
	rel2.Version = 2
	upAction.cfg.Releases.Create(rel2)

	vals := map[string]interface{}{}

	_, err := upAction.Run(rel.Name, buildChart(), vals)
	req.Contains(err.Error(), "progress", err)
}
<|endoftext|>


<commit_before>package agent

import (
	"bufio"
	"context"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/rancher/k3s/pkg/daemons/config"
	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/util/net"
	"k8s.io/component-base/logs"
	app2 "k8s.io/kubernetes/cmd/kube-proxy/app"
	"k8s.io/kubernetes/cmd/kubelet/app"
	_ "k8s.io/kubernetes/pkg/client/metrics/prometheus" // for client metric registration
	_ "k8s.io/kubernetes/pkg/version/prometheus"        // for version metric registration
)

func Agent(config *config.Agent) error {
	rand.Seed(time.Now().UTC().UnixNano())

	kubelet(config)
	kubeProxy(config)

	return nil
}

func kubeProxy(config *config.Agent) {
	args := []string{
		"--proxy-mode", "iptables",
		"--healthz-bind-address", "127.0.0.1",
		"--kubeconfig", config.KubeConfig,
		"--cluster-cidr", config.ClusterCIDR.String(),
	}
	args = append(args, config.ExtraKubeletArgs...)

	command := app2.NewProxyCommand()
	command.SetArgs(args)
	go func() {
		err := command.Execute()
		logrus.Fatalf("kube-proxy exited: %v", err)
	}()
}

func kubelet(cfg *config.Agent) {
	command := app.NewKubeletCommand(context.Background().Done())
	logs.InitLogs()
	defer logs.FlushLogs()

	args := []string{
		"--healthz-bind-address", "127.0.0.1",
		"--read-only-port", "0",
		"--allow-privileged=true",
		"--cluster-domain", "cluster.local",
		"--kubeconfig", cfg.KubeConfig,
		"--eviction-hard", "imagefs.available<5%,nodefs.available<5%",
		"--eviction-minimum-reclaim", "imagefs.available=10%,nodefs.available=10%",
		"--fail-swap-on=false",
		//"--cgroup-root", "/k3s",
		"--cgroup-driver", "cgroupfs",
	}
	if cfg.RootDir != "" {
		args = append(args, "--root-dir", cfg.RootDir)
		args = append(args, "--cert-dir", filepath.Join(cfg.RootDir, "pki"))
		args = append(args, "--seccomp-profile-root", filepath.Join(cfg.RootDir, "seccomp"))
	}
	if cfg.CNIConfDir != "" {
		args = append(args, "--cni-conf-dir", cfg.CNIConfDir)
	}
	if cfg.CNIBinDir != "" {
		args = append(args, "--cni-bin-dir", cfg.CNIBinDir)
	}
	if len(cfg.ClusterDNS) > 0 {
		args = append(args, "--cluster-dns", cfg.ClusterDNS.String())
	}
	if cfg.ResolvConf != "" {
		args = append(args, "--resolv-conf", cfg.ResolvConf)
	}
	if cfg.RuntimeSocket != "" {
		args = append(args, "--container-runtime", "remote")
		args = append(args, "--container-runtime-endpoint", cfg.RuntimeSocket)
	}
	if cfg.ListenAddress != "" {
		args = append(args, "--address", cfg.ListenAddress)
	}
	if cfg.CACertPath != "" {
		args = append(args, "--anonymous-auth=false", "--client-ca-file", cfg.CACertPath)
	}
	if cfg.NodeName != "" {
		args = append(args, "--hostname-override", cfg.NodeName)
	}
	defaultIP, err := net.ChooseHostInterface()
	if err != nil || defaultIP.String() != cfg.NodeIP {
		args = append(args, "--node-ip", cfg.NodeIP)
	}
	root, hasCFS := checkCgroups()
	if !hasCFS {
		logrus.Warn("Disabling CPU quotas due to missing cpu.cfs_period_us")
		args = append(args, "--cpu-cfs-quota=false")
	}
	if root != "" {
		args = append(args, "--runtime-cgroups", root)
		args = append(args, "--kubelet-cgroups", root)
	}
	args = append(args, cfg.ExtraKubeletArgs...)

	command.SetArgs(args)

	go func() {
		logrus.Infof("Running kubelet %s", config.ArgString(args))
		logrus.Fatalf("kubelet exited: %v", command.Execute())
	}()
}

func checkCgroups() (string, bool) {
	f, err := os.Open("/proc/self/cgroup")
	if err != nil {
		return "", false
	}
	defer f.Close()

	ret := false
	root := ""
	scan := bufio.NewScanner(f)
	for scan.Scan() {
		parts := strings.Split(scan.Text(), ":")
		if len(parts) < 3 {
			continue
		}
		systems := strings.Split(parts[1], ",")
		for _, system := range systems {
			if system == "cpu" {
				p := filepath.Join("/sys/fs/cgroup", parts[1], parts[2], "cpu.cfs_period_us")
				if _, err := os.Stat(p); err == nil {
					ret = true
				}
			} else if system == "name=systemd" {
				last := parts[len(parts)-1]
				i := strings.LastIndex(last, ".slice")
				if i > 0 {
					root = "/systemd" + last[:i+len(".slice")]
				}
			}
		}
	}

	return root, ret
}
<commit_msg>Pull in parallel for CRI<commit_after>package agent

import (
	"bufio"
	"context"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/rancher/k3s/pkg/daemons/config"
	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/util/net"
	"k8s.io/component-base/logs"
	app2 "k8s.io/kubernetes/cmd/kube-proxy/app"
	"k8s.io/kubernetes/cmd/kubelet/app"
	_ "k8s.io/kubernetes/pkg/client/metrics/prometheus" // for client metric registration
	_ "k8s.io/kubernetes/pkg/version/prometheus"        // for version metric registration
)

func Agent(config *config.Agent) error {
	rand.Seed(time.Now().UTC().UnixNano())

	kubelet(config)
	kubeProxy(config)

	return nil
}

func kubeProxy(config *config.Agent) {
	args := []string{
		"--proxy-mode", "iptables",
		"--healthz-bind-address", "127.0.0.1",
		"--kubeconfig", config.KubeConfig,
		"--cluster-cidr", config.ClusterCIDR.String(),
	}
	args = append(args, config.ExtraKubeletArgs...)

	command := app2.NewProxyCommand()
	command.SetArgs(args)
	go func() {
		err := command.Execute()
		logrus.Fatalf("kube-proxy exited: %v", err)
	}()
}

func kubelet(cfg *config.Agent) {
	command := app.NewKubeletCommand(context.Background().Done())
	logs.InitLogs()
	defer logs.FlushLogs()

	args := []string{
		"--healthz-bind-address", "127.0.0.1",
		"--read-only-port", "0",
		"--allow-privileged=true",
		"--cluster-domain", "cluster.local",
		"--kubeconfig", cfg.KubeConfig,
		"--eviction-hard", "imagefs.available<5%,nodefs.available<5%",
		"--eviction-minimum-reclaim", "imagefs.available=10%,nodefs.available=10%",
		"--fail-swap-on=false",
		//"--cgroup-root", "/k3s",
		"--cgroup-driver", "cgroupfs",
	}
	if cfg.RootDir != "" {
		args = append(args, "--root-dir", cfg.RootDir)
		args = append(args, "--cert-dir", filepath.Join(cfg.RootDir, "pki"))
		args = append(args, "--seccomp-profile-root", filepath.Join(cfg.RootDir, "seccomp"))
	}
	if cfg.CNIConfDir != "" {
		args = append(args, "--cni-conf-dir", cfg.CNIConfDir)
	}
	if cfg.CNIBinDir != "" {
		args = append(args, "--cni-bin-dir", cfg.CNIBinDir)
	}
	if len(cfg.ClusterDNS) > 0 {
		args = append(args, "--cluster-dns", cfg.ClusterDNS.String())
	}
	if cfg.ResolvConf != "" {
		args = append(args, "--resolv-conf", cfg.ResolvConf)
	}
	if cfg.RuntimeSocket != "" {
		args = append(args, "--container-runtime", "remote")
		args = append(args, "--container-runtime-endpoint", cfg.RuntimeSocket)
		args = append(args, "--serialize-image-pulls=false")
	}
	if cfg.ListenAddress != "" {
		args = append(args, "--address", cfg.ListenAddress)
	}
	if cfg.CACertPath != "" {
		args = append(args, "--anonymous-auth=false", "--client-ca-file", cfg.CACertPath)
	}
	if cfg.NodeName != "" {
		args = append(args, "--hostname-override", cfg.NodeName)
	}
	defaultIP, err := net.ChooseHostInterface()
	if err != nil || defaultIP.String() != cfg.NodeIP {
		args = append(args, "--node-ip", cfg.NodeIP)
	}
	root, hasCFS := checkCgroups()
	if !hasCFS {
		logrus.Warn("Disabling CPU quotas due to missing cpu.cfs_period_us")
		args = append(args, "--cpu-cfs-quota=false")
	}
	if root != "" {
		args = append(args, "--runtime-cgroups", root)
		args = append(args, "--kubelet-cgroups", root)
	}
	args = append(args, cfg.ExtraKubeletArgs...)

	command.SetArgs(args)

	go func() {
		logrus.Infof("Running kubelet %s", config.ArgString(args))
		logrus.Fatalf("kubelet exited: %v", command.Execute())
	}()
}

func checkCgroups() (string, bool) {
	f, err := os.Open("/proc/self/cgroup")
	if err != nil {
		return "", false
	}
	defer f.Close()

	ret := false
	root := ""
	scan := bufio.NewScanner(f)
	for scan.Scan() {
		parts := strings.Split(scan.Text(), ":")
		if len(parts) < 3 {
			continue
		}
		systems := strings.Split(parts[1], ",")
		for _, system := range systems {
			if system == "cpu" {
				p := filepath.Join("/sys/fs/cgroup", parts[1], parts[2], "cpu.cfs_period_us")
				if _, err := os.Stat(p); err == nil {
					ret = true
				}
			} else if system == "name=systemd" {
				last := parts[len(parts)-1]
				i := strings.LastIndex(last, ".slice")
				if i > 0 {
					root = "/systemd" + last[:i+len(".slice")]
				}
			}
		}
	}

	return root, ret
}
<|endoftext|>