Support subset restore with on-error-continue mode (#131)

Add possibility to use on-error-continue with filtering.

Add the discardData method to the IRestoreReader interface. When an error
happens during COPY FROM execution, discardData is called to read the rest of
the table data from plugin's pipe. So the transition to the next table occurs
and the offset is positioned automatically.

In the case of discard error we don't stop the helper, because other helpers can
restore data on other segments. Another reason - hanging on COPY happens when
one helper stops before restoring a next table. Add the discardErr boolean field
to RestoreReader. This field is used in discardData and copyData to avoid
bufReader reading when an error has happened in discardData before.

When skip file is found, call discardData to move on the next table data.

Data written in writeTestTOC has now started to be used. Fix this data.

Ticket: ADBDEV-7890

Author: Andrey Sokolov

helper/helper_test.go

@@ -3,7 +3,9 @@ package helper
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"os"
+	"strings"
 
 	"github.com/greenplum-db/gpbackup/utils"
 	"golang.org/x/sys/unix"
@@ -21,7 +23,9 @@ var (
 )
 
 type restoreReaderTestImpl struct {
-	waitCount int
+	waitCount      int
+	discardedBytes int64
+	discardErr     error
 }
 
 func (r *restoreReaderTestImpl) waitForPlugin() error {
@@ -45,7 +49,16 @@ func (r *restoreReaderTestImpl) closeFileHandle() {
 }
 
 func (r *restoreReaderTestImpl) getReaderType() ReaderType {
-	return "nil"
+	return SUBSET
+}
+
+func (r *restoreReaderTestImpl) discardData(num int64) (int64, error) {
+	if r.discardErr != nil {
+		return 0, r.discardErr
+	}
+
+	r.discardedBytes += num
+	return num, nil
 }
 
 type helperTestStep struct {
@@ -198,6 +211,47 @@ func (pt *testPluginCmd) Wait() error {
 func (pt *testPluginCmd) errLog() {
 }
 
+type limitReader struct {
+	remainder int
+	err       error
+}
+
+func (r *limitReader) Read(p []byte) (n int, err error) {
+	if r.remainder <= 0 {
+		return 0, r.err
+	}
+
+	if len(p) > r.remainder {
+		p = p[0:r.remainder]
+	}
+
+	n = len(p)
+	for i := 0; i < n; i++ {
+		p[i] = 1
+	}
+	r.remainder -= n
+	return
+}
+
+type limitWriter struct {
+	remainder int
+}
+
+func (w *limitWriter) Write(p []byte) (n int, err error) {
+	if w.remainder < len(p) {
+		n = w.remainder
+	} else {
+		n = len(p)
+	}
+
+	if w.remainder == 0 {
+		err = io.ErrShortWrite
+	}
+
+	w.remainder -= n
+	return
+}
+
 var _ = Describe("helper tests", func() {
 	var pluginConfig utils.PluginConfig
 	var isSubset bool
@@ -243,7 +297,7 @@ var _ = Describe("helper tests", func() {
 			pluginConfig.Options["restore_subset"] = "on"
 			*onErrorContinue = true
 			isSubset = getSubsetFlag(fileToRead, &pluginConfig)
-			Expect(isSubset).To(Equal(false))
+			Expect(isSubset).To(Equal(true))
 		})
 		It("when restore_subset is off, --on-error-continue is false, compression \"gz\" is used", func() {
 			pluginConfig.Options["restore_subset"] = "off"
@@ -416,6 +470,52 @@ var _ = Describe("helper tests", func() {
 			err := doRestoreAgentInternal(helper)
 			Expect(err).To(BeNil())
 		})
+		It("discard data if skip file is discovered with single datafile", func() {
+			*singleDataFile = true
+			*isResizeRestore = false
+			*tocFile = testTocFile
+
+			writeTestTOC(testTocFile)
+			defer func() {
+				_ = os.Remove(*tocFile)
+			}()
+
+			oidBatch := []oidWithBatch{
+				{1 /* The first oid from TOC */, 0},
+			}
+
+			expectedScenario := []helperTestStep{
+				{"mock_1_0", false, 1, true, "Can not open pipe for table 1, check_skip_file shall called, skip file exists"},
+			}
+
+			helper := newHelperTest(oidBatch, expectedScenario)
+			err := doRestoreAgentInternal(helper)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(helper.restoreData.discardedBytes).To(Equal(int64(18)))
+		})
+		It("discard error data if skip file is discovered with single datafile", func() {
+			*singleDataFile = true
+			*isResizeRestore = false
+			*tocFile = testTocFile
+
+			writeTestTOC(testTocFile)
+			defer func() {
+				_ = os.Remove(*tocFile)
+			}()
+
+			oidBatch := []oidWithBatch{
+				{1 /* The first oid from TOC */, 0},
+			}
+
+			expectedScenario := []helperTestStep{
+				{"mock_1_0", false, 1, true, "Can not open pipe for table 1, check_skip_file shall called, skip file exists"},
+			}
+
+			helper := newHelperTest(oidBatch, expectedScenario)
+			helper.restoreData.discardErr = io.EOF
+			err := doRestoreAgentInternal(helper)
+			Expect(err).To(Equal(io.EOF))
+		})
 		It("calls Wait in waitForPlugin doRestoreAgent for single data file", func() {
 			*singleDataFile = true
 			*isResizeRestore = false
@@ -476,6 +576,10 @@ var _ = Describe("helper tests", func() {
 		})
 	})
 	Describe("RestoreReader tests", func() {
+		AfterEach(func() {
+			*onErrorContinue = false
+			writer = nil
+		})
 		It("waitForPlugin normal completion", func() {
 			test_cmd1 := testPluginCmd{hasProcess_: true}
 			test_reader := new(RestoreReader)
@@ -504,20 +608,123 @@ var _ = Describe("helper tests", func() {
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(Equal(msg))
 		})
+		It("CopyData, readerType is SUBSET. Normal completion", func() {
+			writer = bufio.NewWriterSize(&limitWriter{100}, 5)
+
+			test_reader := RestoreReader{
+				readerType: SUBSET,
+				bufReader:  bufio.NewReader(&limitReader{100, io.EOF}),
+			}
+
+			bytesRead, err := test_reader.copyData(18)
+			Expect(bytesRead).To(Equal(int64(18)))
+			Expect(err).ToNot(HaveOccurred())
+		})
+		It("CopyData, readerType is SUBSET. Error on write", func() {
+			*onErrorContinue = true
+			bufSize := 5
+			toRead := int64(18)
+			writer = bufio.NewWriterSize(&limitWriter{7}, bufSize)
+
+			test_reader := RestoreReader{
+				readerType: SUBSET,
+				bufReader:  bufio.NewReader(&limitReader{100, io.EOF}),
+			}
+
+			bytesRead, err := test_reader.copyData(toRead)
+			Expect(bytesRead).To(Equal(toRead))
+			Expect(errors.Is(err, io.ErrShortWrite)).To(Equal(true))
+			str := fmt.Sprintf("copied %d bytes from %d: ", bufSize*2, toRead)
+			Expect(strings.HasPrefix(err.Error(), str)).To(Equal(true))
+
+		})
+		It("CopyData, readerType is SUBSET. EOF", func() {
+			*onErrorContinue = true
+			writer = bufio.NewWriterSize(&limitWriter{100}, 5)
+
+			test_reader := RestoreReader{
+				readerType: SUBSET,
+				bufReader:  bufio.NewReader(&limitReader{25, io.EOF}),
+			}
+
+			bytesRead, err := test_reader.copyData(30)
+			Expect(bytesRead).To(Equal(int64(25)))
+			Expect(err).To(Equal(io.EOF))
+		})
+		It("CopyData, readerType is SUBSET. Error on write and EOF", func() {
+			*onErrorContinue = true
+			bufSize := 5
+			toCopy := int64(30)
+			rLmt := int64(25)
+			writer = bufio.NewWriterSize(&limitWriter{7}, bufSize)
+
+			test_reader := RestoreReader{
+				readerType: SUBSET,
+				bufReader:  bufio.NewReader(&limitReader{int(rLmt), io.EOF}),
+			}
+
+			bytesRead, err := test_reader.copyData(toCopy)
+			Expect(bytesRead).To(Equal(rLmt))
+			Expect(errors.Is(err, io.ErrShortWrite)).To(Equal(true))
+			Expect(errors.Is(err, io.EOF)).To(Equal(true))
+			readBeforeErr := int64(bufSize * 2)
+			prefix := fmt.Sprintf("discard error in copyData: discarded %d bytes from %d: ", rLmt-readBeforeErr, toCopy-readBeforeErr)
+			Expect(strings.HasPrefix(err.Error(), prefix)).To(Equal(true))
+			strCopied := fmt.Sprintf("copied %d bytes from %d: ", readBeforeErr, toCopy)
+			Expect(strings.Contains(err.Error(), strCopied)).To(Equal(true))
+
+			bytesRead, err = test_reader.copyData(10)
+			Expect(bytesRead).To(Equal(int64(0)))
+			Expect(err.Error()).To(Equal("10 bytes to copy, but discard error has already occurred. Skipping read."))
+
+			bytesRead, err = test_reader.discardData(5)
+			Expect(bytesRead).To(Equal(int64(0)))
+			Expect(err.Error()).To(Equal("5 bytes to discard, but discard error has already occurred. Skipping read."))
+		})
+		It("CopyData, readerType is SUBSET. Error on write and on read", func() {
+			*onErrorContinue = true
+			bufSize := 5
+			toCopy := int64(30)
+			rLmt := int64(25)
+			writer = bufio.NewWriterSize(&limitWriter{7}, bufSize)
+
+			test_reader := RestoreReader{
+				readerType: SUBSET,
+				bufReader:  bufio.NewReader(&limitReader{int(rLmt), io.ErrNoProgress}),
+			}
+
+			bytesRead, err := test_reader.copyData(toCopy)
+			Expect(bytesRead).To(Equal(rLmt))
+			Expect(errors.Is(err, io.ErrShortWrite)).To(Equal(true))
+			Expect(errors.Is(err, io.ErrNoProgress)).To(Equal(true))
+			readBeforeErr := int64(bufSize * 2)
+			prefix := fmt.Sprintf("discard error in copyData: discarded %d bytes from %d: ", rLmt-readBeforeErr, toCopy-readBeforeErr)
+			Expect(strings.HasPrefix(err.Error(), prefix)).To(Equal(true))
+			strCopied := fmt.Sprintf("copied %d bytes from %d: ", readBeforeErr, toCopy)
+			Expect(strings.Contains(err.Error(), strCopied)).To(Equal(true))
+
+			bytesRead, err = test_reader.copyData(10)
+			Expect(bytesRead).To(Equal(int64(0)))
+			Expect(err.Error()).To(Equal("10 bytes to copy, but discard error has already occurred. Skipping read."))
+
+			bytesRead, err = test_reader.discardData(5)
+			Expect(bytesRead).To(Equal(int64(0)))
+			Expect(err.Error()).To(Equal("5 bytes to discard, but discard error has already occurred. Skipping read."))
+		})
 	})
 })
 
 func writeTestTOC(tocFile string) {
 	// Write test TOC. We are not going to read data using it, so dataLength is a random number
 	dataLength := 100
 	customTOC := fmt.Sprintf(`dataentries:
-1:
+  1:
     startbyte: 0
     endbyte: 18
-2:
+  2:
     startbyte: 18
     endbyte: %[1]d
-3:
+  3:
     startbyte: %[1]d
     endbyte: %d
 `, dataLength+18, dataLength+18+18)

helper/restore_helper.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"compress/gzip"
+	errorsStd "errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -54,6 +55,7 @@ type IRestoreReader interface {
 	copyAllData() (int64, error)
 	closeFileHandle()
 	getReaderType() ReaderType
+	discardData(num int64) (int64, error)
 }
 
 type RestoreReader struct {
@@ -62,6 +64,7 @@ type RestoreReader struct {
 	seekReader io.ReadSeeker
 	pluginCmd  IPluginCmd
 	readerType ReaderType
+	discardErr bool
 }
 
 // Wait for plugin process that should be already finished. This should be
@@ -102,14 +105,53 @@ func (r *RestoreReader) positionReader(pos uint64, oid int) error {
 	return nil
 }
 
+func (r *RestoreReader) discardData(num int64) (int64, error) {
+	if r.readerType != SUBSET {
+		panic("discardData should be called for readerType == SUBSET only")
+	}
+
+	if r.discardErr {
+		err := fmt.Errorf("%d bytes to discard, but discard error has already occurred. Skipping read.", num)
+		logVerbose(err.Error())
+		return 0, err
+	}
+
+	n, err := io.CopyN(io.Discard, r.bufReader, num)
+	if err == nil {
+		logVerbose(fmt.Sprintf("discarded %d bytes", n))
+	} else {
+		r.discardErr = true
+		err = errors.Wrapf(err, "discarded %d bytes from %d", n, num)
+		logError(err.Error())
+	}
+	return n, err
+}
+
 func (r *RestoreReader) copyData(num int64) (int64, error) {
 	var bytesRead int64
 	var err error
 	switch r.readerType {
 	case SEEKABLE:
 		bytesRead, err = io.CopyN(writer, r.seekReader, num)
-	case NONSEEKABLE, SUBSET:
+	case NONSEEKABLE:
+		bytesRead, err = io.CopyN(writer, r.bufReader, num)
+	case SUBSET:
+		if r.discardErr {
+			err := fmt.Errorf("%d bytes to copy, but discard error has already occurred. Skipping read.", num)
+			logVerbose(err.Error())
+			return 0, err
+		}
+
 		bytesRead, err = io.CopyN(writer, r.bufReader, num)
+		if err != nil && err != io.EOF && *onErrorContinue {
+			err = errors.Wrapf(err, "copied %d bytes from %d", bytesRead, num)
+			bytesDiscard, errDiscard := r.discardData(num - bytesRead)
+			bytesRead += bytesDiscard
+			if errDiscard != nil {
+				err = errorsStd.Join(errDiscard, err)
+				err = errors.Wrap(err, "discard error in copyData")
+			}
+		}
 	}
 	return bytesRead, err
 }
@@ -351,6 +393,10 @@ func doRestoreAgentInternal(restoreHelper IRestoreHelper) error {
 						logWarn(fmt.Sprintf("Oid %d, Batch %d: Skip file discovered, skipping this relation.", tableOid, batchNum))
 						err = nil
 						skipOid = tableOid
+						if *singleDataFile && readers[contentToRestore] != nil && readers[contentToRestore].getReaderType() == SUBSET {
+							bytesToDiscard := int64(end[contentToRestore] - start[contentToRestore])
+							_, err = readers[contentToRestore].discardData(bytesToDiscard)
+						}
 						/* Close up to *copyQueue files with this tableOid */
 						for idx := 0; idx < *copyQueue; idx++ {
 							batchToDelete := batchNum + idx
@@ -613,7 +659,7 @@ func getSubsetFlag(fileToRead string, pluginConfig *utils.PluginConfig) bool {
 		return false
 	}
 	// Helper's option does not allow to use subset
-	if !*isFiltered || *onErrorContinue {
+	if !*isFiltered {
 		return false
 	}
 	// Restore subset and compression does not allow together