1// Copyright 2023 The Cockroach Authors.
2//
3// Use of this software is governed by the CockroachDB Software License
4// included in the /LICENSE file.
5
6package nodelocal
7
8import (
9 "context"
10 "io"
11 "net/url"
12
13 "github.com/cockroachdb/cockroach/pkg/base"
14 "github.com/cockroachdb/cockroach/pkg/blobs"
15 "github.com/cockroachdb/cockroach/pkg/cloud"
16 "github.com/cockroachdb/cockroach/pkg/cloud/cloudpb"
17 "github.com/cockroachdb/cockroach/pkg/settings/cluster"
18 "github.com/cockroachdb/cockroach/pkg/util/buildutil"
19 "github.com/cockroachdb/cockroach/pkg/util/ioctx"
20)
21
22// ReplaceNodeLocalForTesting replaces the implementation of of
23// nodelocal with one reads and writes directly from the given
24// directory.
25//
…lines 26–65 unchanged (40 lines)
66 blobClient: c,
67 settings: settings,
68 uri: es.URI,
69 }
70}
71
72// ReplaceNodeLocalForTestingWithInterceptor is like ReplaceNodeLocalForTesting
73// but calls interceptRead before each ReadFile call. The interceptor receives
74// the basename of the file being read. This can be used to inject blocking or
75// delays for specific files, e.g. to block reads of data SSTs during online
76// restore testing. Size calls are not intercepted so that Pebble can still
77// query external file sizes for disk usage accounting.
78func ReplaceNodeLocalForTestingWithInterceptor(
79 root string, interceptRead func(context.Context, string),
80) func() {
81 makeFn := func(ctx context.Context, conf cloud.EarlyBootExternalStorageContext, es cloudpb.ExternalStorage) (cloud.ExternalStorage, error) {
82 inner := TestingMakeNodelocalStorage(root, conf.Settings, es)
83 return &interceptingStorage{inner: inner, interceptRead: interceptRead}, nil
84 }
85 parserFn := func(uri *url.URL) (cloudpb.ExternalStorage, error) {
86 if !buildutil.CrdbTestBuild {
87 panic("nodelocal test implementation in non-test build")
88 }
89 return cloudpb.ExternalStorage{
90 Provider: cloudpb.ExternalStorageProvider_nodelocal,
91 LocalFileConfig: cloudpb.ExternalStorage_LocalFileConfig{
92 Path: uri.Path,
93 },
94 URI: uri.String(),
95 }, nil
96 }
97 return cloud.ReplaceProviderForTesting(cloudpb.ExternalStorageProvider_nodelocal, cloud.RegisteredProvider{
98 EarlyBootConstructFn: makeFn,
99 EarlyBootParseFn: parserFn,
100 Schemes: []string{scheme},
101 })
102}
103
104// interceptingStorage wraps a cloud.ExternalStorage and calls an interceptor
105// before ReadFile operations.
106type interceptingStorage struct {
107 inner cloud.ExternalStorage
108 interceptRead func(context.Context, string)
109}
110
111var _ cloud.ExternalStorage = &interceptingStorage{}
112
113func (s *interceptingStorage) Close() error { return s.inner.Close() }
114func (s *interceptingStorage) Conf() cloudpb.ExternalStorage {
115 return s.inner.Conf()
116}
117func (s *interceptingStorage) ExternalIOConf() base.ExternalIODirConfig {
118 return s.inner.ExternalIOConf()
119}
120func (s *interceptingStorage) RequiresExternalIOAccounting() bool {
121 return s.inner.RequiresExternalIOAccounting()
122}
123func (s *interceptingStorage) Settings() *cluster.Settings {
124 return s.inner.Settings()
125}
126func (s *interceptingStorage) Writer(ctx context.Context, basename string) (io.WriteCloser, error) {
127 return s.inner.Writer(ctx, basename)
128}
129func (s *interceptingStorage) List(
130 ctx context.Context, prefix string, opts cloud.ListOptions, fn cloud.ListingFn,
131) error {
132 return s.inner.List(ctx, prefix, opts, fn)
133}
134func (s *interceptingStorage) Delete(ctx context.Context, basename string) error {
135 return s.inner.Delete(ctx, basename)
136}
137func (s *interceptingStorage) ReadFile(
138 ctx context.Context, basename string, opts cloud.ReadOptions,
139) (ioctx.ReadCloserCtx, int64, error) {
140 s.interceptRead(ctx, basename)
141 return s.inner.ReadFile(ctx, basename, opts)
142}
143func (s *interceptingStorage) Size(ctx context.Context, basename string) (int64, error) {
144 return s.inner.Size(ctx, basename)
145}
146