Optimize ReconstructSome for leopard 8+16

klauspost · Feb 24, 2024 · 0b4e3cb · 0b4e3cb
1 parent 2124328
commit 0b4e3cb
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 21 deletions.
diff --git a/leopard.go b/leopard.go
@@ -334,17 +334,17 @@ func (r *leopardFF16) Split(data []byte) ([][]byte, error) {
 
 func (r *leopardFF16) ReconstructSome(shards [][]byte, required []bool) error {
 	if len(required) == r.totalShards {
-		return r.reconstruct(shards, true)
+		return r.reconstruct(shards, true, required)
 	}
-	return r.reconstruct(shards, false)
+	return r.reconstruct(shards, false, required)
 }
 
 func (r *leopardFF16) Reconstruct(shards [][]byte) error {
-	return r.reconstruct(shards, true)
+	return r.reconstruct(shards, true, nil)
 }
 
 func (r *leopardFF16) ReconstructData(shards [][]byte) error {
-	return r.reconstruct(shards, false)
+	return r.reconstruct(shards, false, nil)
 }
 
 func (r *leopardFF16) Verify(shards [][]byte) (bool, error) {
@@ -375,8 +375,8 @@ func (r *leopardFF16) Verify(shards [][]byte) (bool, error) {
 	return true, nil
 }
 
-func (r *leopardFF16) reconstruct(shards [][]byte, recoverAll bool) error {
-	if len(shards) != r.totalShards {
+func (r *leopardFF16) reconstruct(shards [][]byte, recoverAll bool, required []bool) error {
+	if len(shards) != r.totalShards || required != nil && len(required) < r.dataShards {
 		return ErrTooFewShards
 	}
 
@@ -388,15 +388,19 @@ func (r *leopardFF16) reconstruct(shards [][]byte, recoverAll bool) error {
 	// nothing to do.
 	numberPresent := 0
 	dataPresent := 0
+	missingRequired := 0
 	for i := 0; i < r.totalShards; i++ {
 		if len(shards[i]) != 0 {
 			numberPresent++
 			if i < r.dataShards {
 				dataPresent++
 			}
+		} else if required != nil && required[i] {
+			missingRequired++
 		}
 	}
-	if numberPresent == r.totalShards || !recoverAll && dataPresent == r.dataShards {
+	if numberPresent == r.totalShards || !recoverAll && dataPresent == r.dataShards ||
+		required != nil && missingRequired == 0 {
 		// Cool. All of the shards have data. We don't
 		// need to do anything.
 		return nil
@@ -541,7 +545,7 @@ func (r *leopardFF16) reconstruct(shards [][]byte, recoverAll bool) error {
 		end = r.totalShards
 	}
 	for i := 0; i < end; i++ {
-		if len(shards[i]) != 0 {
+		if len(shards[i]) != 0 || required != nil && !required[i] {
 			continue
 		}
 		if cap(shards[i]) >= shardSize {

diff --git a/leopard8.go b/leopard8.go
@@ -370,17 +370,17 @@ func (r *leopardFF8) Split(data []byte) ([][]byte, error) {
 
 func (r *leopardFF8) ReconstructSome(shards [][]byte, required []bool) error {
 	if len(required) == r.totalShards {
-		return r.reconstruct(shards, true)
+		return r.reconstruct(shards, true, required)
 	}
-	return r.reconstruct(shards, false)
+	return r.reconstruct(shards, false, required)
 }
 
 func (r *leopardFF8) Reconstruct(shards [][]byte) error {
-	return r.reconstruct(shards, true)
+	return r.reconstruct(shards, true, nil)
 }
 
 func (r *leopardFF8) ReconstructData(shards [][]byte) error {
-	return r.reconstruct(shards, false)
+	return r.reconstruct(shards, false, nil)
 }
 
 func (r *leopardFF8) Verify(shards [][]byte) (bool, error) {
@@ -411,8 +411,8 @@ func (r *leopardFF8) Verify(shards [][]byte) (bool, error) {
 	return true, nil
 }
 
-func (r *leopardFF8) reconstruct(shards [][]byte, recoverAll bool) error {
-	if len(shards) != r.totalShards {
+func (r *leopardFF8) reconstruct(shards [][]byte, recoverAll bool, required []bool) error {
+	if len(shards) != r.totalShards || required != nil && len(required) < r.dataShards {
 		return ErrTooFewShards
 	}
 
@@ -424,15 +424,19 @@ func (r *leopardFF8) reconstruct(shards [][]byte, recoverAll bool) error {
 	// nothing to do.
 	numberPresent := 0
 	dataPresent := 0
+	missingRequired := 0
 	for i := 0; i < r.totalShards; i++ {
 		if len(shards[i]) != 0 {
 			numberPresent++
 			if i < r.dataShards {
 				dataPresent++
 			}
+		} else if required != nil && required[i] {
+			missingRequired++
 		}
 	}
-	if numberPresent == r.totalShards || !recoverAll && dataPresent == r.dataShards {
+	if numberPresent == r.totalShards || !recoverAll && dataPresent == r.dataShards ||
+		required != nil && missingRequired == 0 {
 		// Cool. All of the shards have data. We don't
 		// need to do anything.
 		return nil
@@ -566,7 +570,7 @@ func (r *leopardFF8) reconstruct(shards [][]byte, recoverAll bool) error {
 
 	// Add output
 	for i, sh := range shards {
-		if !recoverAll && i >= r.dataShards {
+		if !recoverAll && i >= r.dataShards || required != nil && !required[i] {
 			continue
 		}
 		if len(sh) == 0 {
@@ -657,7 +661,7 @@ func (r *leopardFF8) reconstruct(shards [][]byte, recoverAll bool) error {
 		}
 		// Restore
 		for i := 0; i < end; i++ {
-			if len(sh[i]) != 0 {
+			if len(sh[i]) != 0 || required != nil && !required[i] {
 				continue
 			}
 

diff --git a/reedsolomon_test.go b/reedsolomon_test.go
@@ -877,8 +877,7 @@ func testReconstructData(t *testing.T, o ...Option) {
 	}
 
 	if shardsCopy[2] != nil || shardsCopy[5] != nil || shardsCopy[6] != nil {
-		// This is expected in some cases.
-		t.Log("ReconstructSome reconstructed extra shards")
+		t.Error("ReconstructSome reconstructed extra shards")
 	}
 
 	// Reconstruct with 10 shards present. Use pre-allocated memory for one of them.
@@ -1490,12 +1489,12 @@ func BenchmarkReconstruct10x4x1M(b *testing.B) {
 	benchmarkReconstruct(b, 10, 4, 1024*1024)
 }
 
-// Benchmark 5 data slices with 2 parity slices holding 1MB bytes each
+// Benchmark 50 data slices with 20 parity slices holding 1MB bytes each
 func BenchmarkReconstruct50x20x1M(b *testing.B) {
 	benchmarkReconstruct(b, 50, 20, 1024*1024)
 }
 
-// Benchmark 5 data slices with 2 parity slices holding 1MB bytes each
+// Benchmark 50 data slices with 20 parity slices holding 1MB bytes each
 func BenchmarkReconstructLeopard50x20x1M(b *testing.B) {
 	benchmarkReconstruct(b, 50, 20, 1024*1024, WithLeopardGF(true), WithInversionCache(true))
 }