Repository: pingcap/tidb
File: pkg/objstore/objectio/writer_test.go (Go)

Code before:
// Copyright 2020 PingCAP, 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 objectio_test

import (
	"bytes"
	"context"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"testing/synctest"

	"github.com/klauspost/compress/zstd"
	"github.com/pingcap/tidb/pkg/objstore"
	"github.com/pingcap/tidb/pkg/objstore/compressedio"
	"github.com/pingcap/tidb/pkg/objstore/storeapi"
	"github.com/stretchr/testify/require"
)

func getStore(t *testing.T, uri string, changeStoreFn func(s storeapi.Storage) storeapi.Storage) storeapi.Storage {
	t.Helper()
	backend, err := objstore.ParseBackend(uri, nil)
	require.NoError(t, err)
	ctx := context.Background()
	storage, err := objstore.Create(ctx, backend, true)
	require.NoError(t, err)
	return changeStoreFn(storage)
}

func writeFile(t *testing.T, storage storeapi.Storage, fileName string, lines []string) {
	t.Helper()
	ctx := context.Background()
	writer, err := storage.Create(ctx, fileName, nil)
	require.NoError(t, err)
	for _, str := range lines {
		p := []byte(str)
		written, err2 := writer.Write(ctx, p)
		require.Nil(t, err2)
		require.Len(t, p, written)
	}
	err = writer.Close(ctx)
	require.NoError(t, err)
}

func TestExternalFileWriter(t *testing.T) {
	dir := t.TempDir()

	type testcase struct {
		name    string
		content []string
	}
	testFn := func(test *testcase, t *testing.T) {
		t.Log(test.name)
		storage := getStore(t, "local://"+filepath.ToSlash(dir), func(s storeapi.Storage) storeapi.Storage {
			return s
		})
		fileName := strings.ReplaceAll(test.name, " ", "-") + ".txt"
		writeFile(t, storage, fileName, test.content)
		content, err := os.ReadFile(filepath.Join(dir, fileName))
		require.NoError(t, err)
		require.Equal(t, strings.Join(test.content, ""), string(content))
	}
	tests := []testcase{
		{
			name:    "short and sweet",
			content: []string{"hi"},
		},
		{
			name: "long text small chunks",
			content: []string{
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
			},
		},
		{
			name: "long text medium chunks",
			content: []string{
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
			},
		},
		{
			name: "long text large chunks",
			content: []string{
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
			},
		},
	}
	for i := range tests {
		testFn(&tests[i], t)
	}
}

func TestCompressReaderWriter(t *testing.T) {
	dir := t.TempDir()

	type testcase struct {
		name         string
		content      []string
		compressType compressedio.CompressType
	}
	testFn := func(test *testcase, t *testing.T) {
		t.Log(test.name)
		suffix := test.compressType.FileSuffix()
		fileName := strings.ReplaceAll(test.name, " ", "-") + suffix
		storage := getStore(t, "local://"+filepath.ToSlash(dir), func(s storeapi.Storage) storeapi.Storage {
			return objstore.WithCompression(s, test.compressType, compressedio.DecompressConfig{})
		})
		writeFile(t, storage, fileName, test.content)

		// make sure compressed file is written correctly
		file, err := os.Open(filepath.Join(dir, fileName))
		require.NoError(t, err)
		r, err := compressedio.NewReader(test.compressType, compressedio.DecompressConfig{}, file)
		require.NoError(t, err)
		var bf bytes.Buffer
		_, err = bf.ReadFrom(r)
		require.NoError(t, err)
		require.Equal(t, strings.Join(test.content, ""), bf.String())

		// test withCompression Open
		ctx := context.Background()
		r, err = storage.Open(ctx, fileName, nil)
		require.NoError(t, err)
		content, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, strings.Join(test.content, ""), string(content))

		require.Nil(t, file.Close())
	}
	compressTypeArr := []compressedio.CompressType{compressedio.Gzip, compressedio.Snappy, compressedio.Zstd}

	tests := []testcase{
		{
			name: "long text medium chunks",
			content: []string{
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
			},
		},
		{
			name: "long text large chunks",
			content: []string{
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
				"hello world",
			},
		},
	}
	for i := range tests {
		for _, compressType := range compressTypeArr {
			tests[i].compressType = compressType
			testFn(&tests[i], t)
		}
	}
}

func TestNewCompressReader(t *testing.T) {
	var buf bytes.Buffer
	var w io.WriteCloser
	var err error
	w, err = zstd.NewWriter(&buf)
	require.NoError(t, err)
	_, err = w.Write([]byte("data"))
	require.NoError(t, err)
	require.NoError(t, w.Close())
	compressedData := buf.Bytes()

	// default cfg(decode asynchronously)
	synctest.Test(t, func(t *testing.T) {
		pipeReader, pipeWriter := io.Pipe()

		prevRoutineCnt := runtime.NumGoroutine()
		r, err := compressedio.NewReader(compressedio.Zstd, compressedio.DecompressConfig{}, pipeReader)
		synctest.Wait()

		currRoutineCnt := runtime.NumGoroutine()
		require.NoError(t, err)
		require.Greater(t, currRoutineCnt, prevRoutineCnt)

		_, err = pipeWriter.Write(compressedData)
		require.NoError(t, err)
		err = pipeWriter.Close()
		require.NoError(t, err)

		allData, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, "data", string(allData))
	})

	// sync decode
	synctest.Test(t, func(t *testing.T) {
		pipeReader, pipeWriter := io.Pipe()

		prevRoutineCnt := runtime.NumGoroutine()
		config := compressedio.DecompressConfig{ZStdDecodeConcurrency: 1}
		r, err := compressedio.NewReader(compressedio.Zstd, config, pipeReader)
		synctest.Wait()

		require.NoError(t, err)
		currRoutineCnt := runtime.NumGoroutine()
		require.Equal(t, prevRoutineCnt, currRoutineCnt)

		// because this is "sync" decode, the Write will block and cause a deadlock. We
		// must use another goroutine. This is another sign that the decode is sync and
		// previous synctest is async.
		go func() {
			_, err := pipeWriter.Write(compressedData)
			require.NoError(t, err)
			err = pipeWriter.Close()
			require.NoError(t, err)
		}()

		allData, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, "data", string(allData))
	})
}


Review comment (bug):
Similar to the first test case, if the test fails before pipeWriter is closed (e.g., at lines 235, 236, or 237), the pipe writer will not be properly closed. This is particularly problematic in this test case because the goroutine at line 242 will hang indefinitely trying to write to a pipe that has no reader, if the test exits early. Consider using defer pipeWriter.Close() or t.Cleanup() to ensure proper resource cleanup.

Code after:
		pipeReader, pipeWriter := io.Pipe()
		defer pipeReader.Close()

		prevRoutineCnt := runtime.NumGoroutine()
		config := compressedio.DecompressConfig{ZStdDecodeConcurrency: 1}
		r, err := compressedio.NewReader(compressedio.Zstd, config, pipeReader)
		synctest.Wait()

		require.NoError(t, err)
		currRoutineCnt := runtime.NumGoroutine()
		require.Equal(t, prevRoutineCnt, currRoutineCnt)

		// because this is "sync" decode, the Write will block and cause a deadlock. We
		// must use another goroutine. This is another sign that the decode is sync and
		// previous synctest is async.
		go func() {
			_, err := pipeWriter.Write(compressedData)
			assert.NoError(t, err)
			err = pipeWriter.Close()
			assert.NoError(t, err)
		}()

		allData, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, "data", string(allData))
	})
}