Check in gVisor.

PiperOrigin-RevId: 194583126
Change-Id: Ica1d8821a90f74e7e745962d71801c598c652463

3 files changed, 658 insertions, 0 deletions

diff --git a/pkg/segment/test/BUILD b/pkg/segment/test/BUILD
new file mode 100644
index 000000000..9d398d71a
--- /dev/null
+++ b/pkg/segment/test/BUILD
@@ -0,0 +1,52 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+package(
+    default_visibility = ["//visibility:private"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+load("//tools/go_generics:defs.bzl", "go_template_instance")
+
+go_template_instance(
+    name = "int_range",
+    out = "int_range.go",
+    package = "segment",
+    template = "//pkg/segment:generic_range",
+    types = {
+        "T": "int",
+    },
+)
+
+go_template_instance(
+    name = "int_set",
+    out = "int_set.go",
+    package = "segment",
+    template = "//pkg/segment:generic_set",
+    types = {
+        "Key": "int",
+        "Range": "Range",
+        "Value": "int",
+        "Functions": "setFunctions",
+    },
+)
+
+go_library(
+    name = "segment",
+    testonly = 1,
+    srcs = [
+        "int_range.go",
+        "int_set.go",
+        "set_functions.go",
+    ],
+    importpath = "gvisor.googlesource.com/gvisor/pkg/segment/segment",
+    deps = [
+        "//pkg/state",
+    ],
+)
+
+go_test(
+    name = "segment_test",
+    size = "small",
+    srcs = ["segment_test.go"],
+    embed = [":segment"],
+)
diff --git a/pkg/segment/test/segment_test.go b/pkg/segment/test/segment_test.go
new file mode 100644
index 000000000..7ea24b177
--- /dev/null
+++ b/pkg/segment/test/segment_test.go
@@ -0,0 +1,564 @@
+// Copyright 2018 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package segment
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+)
+
+const (
+	// testSize is the baseline number of elements inserted into sets under
+	// test, and is chosen to be large enough to ensure interesting amounts of
+	// tree rebalancing.
+	//
+	// Note that because checkSet is called between each insertion/removal in
+	// some tests that use it, tests may be quadratic in testSize.
+	testSize = 8000
+
+	// valueOffset is the difference between the value and start of test
+	// segments.
+	valueOffset = 100000
+)
+
+func shuffle(xs []int) {
+	for i := range xs {
+		j := rand.Intn(i + 1)
+		xs[i], xs[j] = xs[j], xs[i]
+	}
+}
+
+func randPermutation(size int) []int {
+	p := make([]int, size)
+	for i := range p {
+		p[i] = i
+	}
+	shuffle(p)
+	return p
+}
+
+// checkSet returns an error if s is incorrectly sorted, does not contain
+// exactly expectedSegments segments, or contains a segment for which val !=
+// key + valueOffset.
+func checkSet(s *Set, expectedSegments int) error {
+	havePrev := false
+	prev := 0
+	nrSegments := 0
+	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+		next := seg.Start()
+		if havePrev && prev >= next {
+			return fmt.Errorf("incorrect order: key %d (segment %d) >= key %d (segment %d)", prev, nrSegments-1, next, nrSegments)
+		}
+		if got, want := seg.Value(), seg.Start()+valueOffset; got != want {
+			return fmt.Errorf("segment %d has key %d, value %d (expected %d)", nrSegments, seg.Start, got, want)
+		}
+		prev = next
+		havePrev = true
+		nrSegments++
+	}
+	if nrSegments != expectedSegments {
+		return fmt.Errorf("incorrect number of segments: got %d, wanted %d", nrSegments, expectedSegments)
+	}
+	return nil
+}
+
+// countSegmentsIn returns the number of segments in s.
+func countSegmentsIn(s *Set) int {
+	var count int
+	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+		count++
+	}
+	return count
+}
+
+func TestAddRandom(t *testing.T) {
+	var s Set
+	order := randPermutation(testSize)
+	var nrInsertions int
+	for i, j := range order {
+		if !s.AddWithoutMerging(Range{j, j + 1}, j+valueOffset) {
+			t.Errorf("Iteration %d: failed to insert segment with key %d", i, j)
+			break
+		}
+		nrInsertions++
+		if err := checkSet(&s, nrInsertions); err != nil {
+			t.Errorf("Iteration %d: %v", i, err)
+			break
+		}
+	}
+	if got, want := countSegmentsIn(&s), nrInsertions; got != want {
+		t.Errorf("Wrong final number of segments: got %d, wanted %d", got, want)
+	}
+	if t.Failed() {
+		t.Logf("Insertion order: %v", order[:nrInsertions])
+		t.Logf("Set contents:\n%v", &s)
+	}
+}
+
+func TestRemoveRandom(t *testing.T) {
+	var s Set
+	for i := 0; i < testSize; i++ {
+		if !s.AddWithoutMerging(Range{i, i + 1}, i+valueOffset) {
+			t.Fatalf("Failed to insert segment %d", i)
+		}
+	}
+	order := randPermutation(testSize)
+	var nrRemovals int
+	for i, j := range order {
+		seg := s.FindSegment(j)
+		if !seg.Ok() {
+			t.Errorf("Iteration %d: failed to find segment with key %d", i, j)
+			break
+		}
+		s.Remove(seg)
+		nrRemovals++
+		if err := checkSet(&s, testSize-nrRemovals); err != nil {
+			t.Errorf("Iteration %d: %v", i, err)
+			break
+		}
+	}
+	if got, want := countSegmentsIn(&s), testSize-nrRemovals; got != want {
+		t.Errorf("Wrong final number of segments: got %d, wanted %d", got, want)
+	}
+	if t.Failed() {
+		t.Logf("Removal order: %v", order[:nrRemovals])
+		t.Logf("Set contents:\n%v", &s)
+		t.FailNow()
+	}
+}
+
+func TestAddSequentialAdjacent(t *testing.T) {
+	var s Set
+	var nrInsertions int
+	for i := 0; i < testSize; i++ {
+		if !s.AddWithoutMerging(Range{i, i + 1}, i+valueOffset) {
+			t.Fatalf("Failed to insert segment %d", i)
+		}
+		nrInsertions++
+		if err := checkSet(&s, nrInsertions); err != nil {
+			t.Errorf("Iteration %d: %v", i, err)
+			break
+		}
+	}
+	if got, want := countSegmentsIn(&s), nrInsertions; got != want {
+		t.Errorf("Wrong final number of segments: got %d, wanted %d", got, want)
+	}
+	if t.Failed() {
+		t.Logf("Set contents:\n%v", &s)
+	}
+
+	first := s.FirstSegment()
+	gotSeg, gotGap := first.PrevNonEmpty()
+	if wantGap := s.FirstGap(); gotSeg.Ok() || gotGap != wantGap {
+		t.Errorf("FirstSegment().PrevNonEmpty(): got (%v, %v), wanted (<terminal iterator>, %v)", gotSeg, gotGap, wantGap)
+	}
+	gotSeg, gotGap = first.NextNonEmpty()
+	if wantSeg := first.NextSegment(); gotSeg != wantSeg || gotGap.Ok() {
+		t.Errorf("FirstSegment().NextNonEmpty(): got (%v, %v), wanted (%v, <terminal iterator>)", gotSeg, gotGap, wantSeg)
+	}
+
+	last := s.LastSegment()
+	gotSeg, gotGap = last.PrevNonEmpty()
+	if wantSeg := last.PrevSegment(); gotSeg != wantSeg || gotGap.Ok() {
+		t.Errorf("LastSegment().PrevNonEmpty(): got (%v, %v), wanted (%v, <terminal iterator>)", gotSeg, gotGap, wantSeg)
+	}
+	gotSeg, gotGap = last.NextNonEmpty()
+	if wantGap := s.LastGap(); gotSeg.Ok() || gotGap != wantGap {
+		t.Errorf("LastSegment().NextNonEmpty(): got (%v, %v), wanted (<terminal iterator>, %v)", gotSeg, gotGap, wantGap)
+	}
+
+	for seg := first.NextSegment(); seg != last; seg = seg.NextSegment() {
+		gotSeg, gotGap = seg.PrevNonEmpty()
+		if wantSeg := seg.PrevSegment(); gotSeg != wantSeg || gotGap.Ok() {
+			t.Errorf("%v.PrevNonEmpty(): got (%v, %v), wanted (%v, <terminal iterator>)", seg, gotSeg, gotGap, wantSeg)
+		}
+		gotSeg, gotGap = seg.NextNonEmpty()
+		if wantSeg := seg.NextSegment(); gotSeg != wantSeg || gotGap.Ok() {
+			t.Errorf("%v.NextNonEmpty(): got (%v, %v), wanted (%v, <terminal iterator>)", seg, gotSeg, gotGap, wantSeg)
+		}
+	}
+}
+
+func TestAddSequentialNonAdjacent(t *testing.T) {
+	var s Set
+	var nrInsertions int
+	for i := 0; i < testSize; i++ {
+		// The range here differs from TestAddSequentialAdjacent so that
+		// consecutive segments are not adjacent.
+		if !s.AddWithoutMerging(Range{2 * i, 2*i + 1}, 2*i+valueOffset) {
+			t.Fatalf("Failed to insert segment %d", i)
+		}
+		nrInsertions++
+		if err := checkSet(&s, nrInsertions); err != nil {
+			t.Errorf("Iteration %d: %v", i, err)
+			break
+		}
+	}
+	if got, want := countSegmentsIn(&s), nrInsertions; got != want {
+		t.Errorf("Wrong final number of segments: got %d, wanted %d", got, want)
+	}
+	if t.Failed() {
+		t.Logf("Set contents:\n%v", &s)
+	}
+
+	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+		gotSeg, gotGap := seg.PrevNonEmpty()
+		if wantGap := seg.PrevGap(); gotSeg.Ok() || gotGap != wantGap {
+			t.Errorf("%v.PrevNonEmpty(): got (%v, %v), wanted (<terminal iterator>, %v)", seg, gotSeg, gotGap, wantGap)
+		}
+		gotSeg, gotGap = seg.NextNonEmpty()
+		if wantGap := seg.NextGap(); gotSeg.Ok() || gotGap != wantGap {
+			t.Errorf("%v.NextNonEmpty(): got (%v, %v), wanted (<terminal iterator>, %v)", seg, gotSeg, gotGap, wantGap)
+		}
+	}
+}
+
+func TestMergeSplit(t *testing.T) {
+	tests := []struct {
+		name      string
+		initial   []Range
+		split     bool
+		splitAddr int
+		final     []Range
+	}{
+		{
+			name:    "Add merges after existing segment",
+			initial: []Range{{1000, 1100}, {1100, 1200}},
+			final:   []Range{{1000, 1200}},
+		},
+		{
+			name:    "Add merges before existing segment",
+			initial: []Range{{1100, 1200}, {1000, 1100}},
+			final:   []Range{{1000, 1200}},
+		},
+		{
+			name:    "Add merges between existing segments",
+			initial: []Range{{1000, 1100}, {1200, 1300}, {1100, 1200}},
+			final:   []Range{{1000, 1300}},
+		},
+		{
+			name:      "SplitAt does nothing at a free address",
+			initial:   []Range{{100, 200}},
+			split:     true,
+			splitAddr: 300,
+			final:     []Range{{100, 200}},
+		},
+		{
+			name:      "SplitAt does nothing at the beginning of a segment",
+			initial:   []Range{{100, 200}},
+			split:     true,
+			splitAddr: 100,
+			final:     []Range{{100, 200}},
+		},
+		{
+			name:      "SplitAt does nothing at the end of a segment",
+			initial:   []Range{{100, 200}},
+			split:     true,
+			splitAddr: 200,
+			final:     []Range{{100, 200}},
+		},
+		{
+			name:      "SplitAt splits in the middle of a segment",
+			initial:   []Range{{100, 200}},
+			split:     true,
+			splitAddr: 150,
+			final:     []Range{{100, 150}, {150, 200}},
+		},
+	}
+Tests:
+	for _, test := range tests {
+		var s Set
+		for _, r := range test.initial {
+			if !s.Add(r, 0) {
+				t.Errorf("%s: Add(%v) failed; set contents:\n%v", test.name, r, &s)
+				continue Tests
+			}
+		}
+		if test.split {
+			s.SplitAt(test.splitAddr)
+		}
+		var i int
+		for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+			if i > len(test.final) {
+				t.Errorf("%s: Incorrect number of segments: got %d, wanted %d; set contents:\n%v", test.name, countSegmentsIn(&s), len(test.final), &s)
+				continue Tests
+			}
+			if got, want := seg.Range(), test.final[i]; got != want {
+				t.Errorf("%s: Segment %d mismatch: got %v, wanted %v; set contents:\n%v", test.name, i, got, want, &s)
+				continue Tests
+			}
+			i++
+		}
+		if i < len(test.final) {
+			t.Errorf("%s: Incorrect number of segments: got %d, wanted %d; set contents:\n%v", test.name, i, len(test.final), &s)
+		}
+	}
+}
+
+func TestIsolate(t *testing.T) {
+	tests := []struct {
+		name    string
+		initial Range
+		bounds  Range
+		final   []Range
+	}{
+		{
+			name:    "Isolate does not split a segment that falls inside bounds",
+			initial: Range{100, 200},
+			bounds:  Range{100, 200},
+			final:   []Range{{100, 200}},
+		},
+		{
+			name:    "Isolate splits at beginning of segment",
+			initial: Range{50, 200},
+			bounds:  Range{100, 200},
+			final:   []Range{{50, 100}, {100, 200}},
+		},
+		{
+			name:    "Isolate splits at end of segment",
+			initial: Range{100, 250},
+			bounds:  Range{100, 200},
+			final:   []Range{{100, 200}, {200, 250}},
+		},
+		{
+			name:    "Isolate splits at beginning and end of segment",
+			initial: Range{50, 250},
+			bounds:  Range{100, 200},
+			final:   []Range{{50, 100}, {100, 200}, {200, 250}},
+		},
+	}
+Tests:
+	for _, test := range tests {
+		var s Set
+		seg := s.Insert(s.FirstGap(), test.initial, 0)
+		seg = s.Isolate(seg, test.bounds)
+		if !test.bounds.IsSupersetOf(seg.Range()) {
+			t.Errorf("%s: Isolated segment %v lies outside bounds %v; set contents:\n%v", test.name, seg.Range(), test.bounds, &s)
+		}
+		var i int
+		for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+			if i > len(test.final) {
+				t.Errorf("%s: Incorrect number of segments: got %d, wanted %d; set contents:\n%v", test.name, countSegmentsIn(&s), len(test.final), &s)
+				continue Tests
+			}
+			if got, want := seg.Range(), test.final[i]; got != want {
+				t.Errorf("%s: Segment %d mismatch: got %v, wanted %v; set contents:\n%v", test.name, i, got, want, &s)
+				continue Tests
+			}
+			i++
+		}
+		if i < len(test.final) {
+			t.Errorf("%s: Incorrect number of segments: got %d, wanted %d; set contents:\n%v", test.name, i, len(test.final), &s)
+		}
+	}
+}
+
+func benchmarkAddSequential(b *testing.B, size int) {
+	for n := 0; n < b.N; n++ {
+		var s Set
+		for i := 0; i < size; i++ {
+			if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+				b.Fatalf("Failed to insert segment %d", i)
+			}
+		}
+	}
+}
+
+func benchmarkAddRandom(b *testing.B, size int) {
+	order := randPermutation(size)
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		var s Set
+		for _, i := range order {
+			if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+				b.Fatalf("Failed to insert segment %d", i)
+			}
+		}
+	}
+}
+
+func benchmarkFindSequential(b *testing.B, size int) {
+	var s Set
+	for i := 0; i < size; i++ {
+		if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+			b.Fatalf("Failed to insert segment %d", i)
+		}
+	}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		for i := 0; i < size; i++ {
+			if seg := s.FindSegment(i); !seg.Ok() {
+				b.Fatalf("Failed to find segment %d", i)
+			}
+		}
+	}
+}
+
+func benchmarkFindRandom(b *testing.B, size int) {
+	var s Set
+	for i := 0; i < size; i++ {
+		if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+			b.Fatalf("Failed to insert segment %d", i)
+		}
+	}
+	order := randPermutation(size)
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		for _, i := range order {
+			if si := s.FindSegment(i); !si.Ok() {
+				b.Fatalf("Failed to find segment %d", i)
+			}
+		}
+	}
+}
+
+func benchmarkIteration(b *testing.B, size int) {
+	var s Set
+	for i := 0; i < size; i++ {
+		if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+			b.Fatalf("Failed to insert segment %d", i)
+		}
+	}
+
+	b.ResetTimer()
+	var count uint64
+	for n := 0; n < b.N; n++ {
+		for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+			count++
+		}
+	}
+	if got, want := count, uint64(size)*uint64(b.N); got != want {
+		b.Fatalf("Iterated wrong number of segments: got %d, wanted %d", got, want)
+	}
+}
+
+func benchmarkAddFindRemoveSequential(b *testing.B, size int) {
+	for n := 0; n < b.N; n++ {
+		var s Set
+		for i := 0; i < size; i++ {
+			if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+				b.Fatalf("Failed to insert segment %d", i)
+			}
+		}
+		for i := 0; i < size; i++ {
+			seg := s.FindSegment(i)
+			if !seg.Ok() {
+				b.Fatalf("Failed to find segment %d", i)
+			}
+			s.Remove(seg)
+		}
+		if !s.IsEmpty() {
+			b.Fatalf("Set not empty after all removals:\n%v", &s)
+		}
+	}
+}
+
+func benchmarkAddFindRemoveRandom(b *testing.B, size int) {
+	order := randPermutation(size)
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		var s Set
+		for _, i := range order {
+			if !s.AddWithoutMerging(Range{i, i + 1}, i) {
+				b.Fatalf("Failed to insert segment %d", i)
+			}
+		}
+		for _, i := range order {
+			seg := s.FindSegment(i)
+			if !seg.Ok() {
+				b.Fatalf("Failed to find segment %d", i)
+			}
+			s.Remove(seg)
+		}
+		if !s.IsEmpty() {
+			b.Fatalf("Set not empty after all removals:\n%v", &s)
+		}
+	}
+}
+
+// Although we don't generally expect our segment sets to get this big, they're
+// useful for emulating the effect of cache pressure.
+var testSizes = []struct {
+	desc string
+	size int
+}{
+	{"64", 1 << 6},
+	{"256", 1 << 8},
+	{"1K", 1 << 10},
+	{"4K", 1 << 12},
+	{"16K", 1 << 14},
+	{"64K", 1 << 16},
+}
+
+func BenchmarkAddSequential(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkAddSequential(b, test.size)
+		})
+	}
+}
+
+func BenchmarkAddRandom(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkAddRandom(b, test.size)
+		})
+	}
+}
+
+func BenchmarkFindSequential(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkFindSequential(b, test.size)
+		})
+	}
+}
+
+func BenchmarkFindRandom(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkFindRandom(b, test.size)
+		})
+	}
+}
+
+func BenchmarkIteration(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkIteration(b, test.size)
+		})
+	}
+}
+
+func BenchmarkAddFindRemoveSequential(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkAddFindRemoveSequential(b, test.size)
+		})
+	}
+}
+
+func BenchmarkAddFindRemoveRandom(b *testing.B) {
+	for _, test := range testSizes {
+		b.Run(test.desc, func(b *testing.B) {
+			benchmarkAddFindRemoveRandom(b, test.size)
+		})
+	}
+}
diff --git a/pkg/segment/test/set_functions.go b/pkg/segment/test/set_functions.go
new file mode 100644
index 000000000..37c196ea1
--- /dev/null
+++ b/pkg/segment/test/set_functions.go
@@ -0,0 +1,42 @@
+// Copyright 2018 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package segment
+
+// Basic numeric constants that we define because the math package doesn't.
+// TODO: These should be Math.MaxInt64/MinInt64?
+const (
+	maxInt = int(^uint(0) >> 1)
+	minInt = -maxInt - 1
+)
+
+type setFunctions struct{}
+
+func (setFunctions) MinKey() int {
+	return minInt
+}
+
+func (setFunctions) MaxKey() int {
+	return maxInt
+}
+
+func (setFunctions) ClearValue(*int) {}
+
+func (setFunctions) Merge(_ Range, val1 int, _ Range, _ int) (int, bool) {
+	return val1, true
+}
+
+func (setFunctions) Split(_ Range, val int, _ int) (int, int) {
+	return val, val
+}