From c3728e2ec9791148fe4315bae630226f419b747e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Leszczy=C5=84ski?= <2000michal@wp.pl>
Date: Wed, 19 Jun 2024 16:24:46 +0200
Subject: [PATCH] refactor(backup): check safety of desc schema in parallel

---
 pkg/scyllaclient/client_agent.go    | 23 +++++++++++++
 pkg/service/backup/worker_schema.go | 51 +++++++++++++++--------------
 2 files changed, 49 insertions(+), 25 deletions(-)

diff --git a/pkg/scyllaclient/client_agent.go b/pkg/scyllaclient/client_agent.go
index 99f494bd0..df036c67e 100644
--- a/pkg/scyllaclient/client_agent.go
+++ b/pkg/scyllaclient/client_agent.go
@@ -180,6 +180,29 @@ func (ni *NodeInfo) SupportsRepairSmallTableOptimization() (bool, error) {
 	return supports, nil
 }
 
+// SupportsSafeDescribeSchemaWithInternals returns true if the output of DESCRIBE SCHEMA WITH INTERNALS
+// is safe to use with backup/restore procedure.
+func (ni *NodeInfo) SupportsSafeDescribeSchemaWithInternals() (bool, error) {
+	// Detect master builds
+	if scyllaversion.MasterVersion(ni.ScyllaVersion) {
+		return true, nil
+	}
+	// Check OSS
+	supports, err := scyllaversion.CheckConstraint(ni.ScyllaVersion, ">= 6.0, < 2000")
+	if err != nil {
+		return false, errors.Errorf("Unsupported Scylla version: %s", ni.ScyllaVersion)
+	}
+	if supports {
+		return true, nil
+	}
+	// Check ENT
+	supports, err = scyllaversion.CheckConstraint(ni.ScyllaVersion, ">= 2024.2, > 1000")
+	if err != nil {
+		return false, errors.Errorf("Unsupported Scylla version: %s", ni.ScyllaVersion)
+	}
+	return supports, nil
+}
+
 // FreeOSMemory calls debug.FreeOSMemory on the agent to return memory to OS.
 func (c *Client) FreeOSMemory(ctx context.Context, host string) error {
 	p := operations.FreeOSMemoryParams{
diff --git a/pkg/service/backup/worker_schema.go b/pkg/service/backup/worker_schema.go
index 512f0beff..e026dc737 100644
--- a/pkg/service/backup/worker_schema.go
+++ b/pkg/service/backup/worker_schema.go
@@ -7,6 +7,7 @@ import (
 	"compress/gzip"
 	"context"
 	"encoding/json"
+	"sync/atomic"
 	"time"
 
 	"github.com/pkg/errors"
@@ -17,7 +18,7 @@ import (
 	"github.com/scylladb/scylla-manager/v3/pkg/util/parallel"
 	"github.com/scylladb/scylla-manager/v3/pkg/util/query"
 	"github.com/scylladb/scylla-manager/v3/pkg/util/retry"
-	"github.com/scylladb/scylla-manager/v3/pkg/util/version"
+	"golang.org/x/sync/errgroup"
 )
 
 func (w *workerTools) AwaitSchemaAgreement(ctx context.Context, clusterSession gocqlx.Session) error {
@@ -146,30 +147,30 @@ func marshalAndCompressSchema(schema query.DescribedSchema) (bytes.Buffer, error
 
 // isDescribeSchemaSafe checks if restoring schema from DESCRIBE SCHEMA WITH INTERNALS is safe.
 func isDescribeSchemaSafe(ctx context.Context, client *scyllaclient.Client, hosts []string) (bool, error) {
-	const (
-		ConstraintOSS = ">= 6.0, < 2000"
-		ConstraintENT = ">= 2024.2, > 1000"
-	)
-
-	for _, h := range hosts {
-		ni, err := client.NodeInfo(ctx, h)
-		if err != nil {
-			return false, errors.Wrapf(err, "get node %s info", h)
-		}
-
-		oss, err := version.CheckConstraint(ni.ScyllaVersion, ConstraintOSS)
-		if err != nil {
-			return false, errors.Wrapf(err, "check version constraint for %s", h)
-		}
-		ent, err := version.CheckConstraint(ni.ScyllaVersion, ConstraintENT)
-		if err != nil {
-			return false, errors.Wrapf(err, "check version constraint for %s", h)
-		}
-
-		if !oss && !ent {
-			return false, nil
-		}
+	out := atomic.Bool{}
+	out.Store(true)
+	eg := errgroup.Group{}
+
+	for _, host := range hosts {
+		h := host
+		eg.Go(func() error {
+			ni, err := client.NodeInfo(ctx, h)
+			if err != nil {
+				return err
+			}
+			res, err := ni.SupportsSafeDescribeSchemaWithInternals()
+			if err != nil {
+				return err
+			}
+			if !res {
+				out.Store(false)
+			}
+			return nil
+		})
 	}
 
-	return true, nil
+	if err := eg.Wait(); err != nil {
+		return false, err
+	}
+	return out.Load(), nil
 }