Implement method to return bucket type from server

go.mod

@@ -71,6 +71,7 @@ require (
 	github.com/rogpeppe/go-internal v1.12.0 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.51.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 // indirect
 	go.opentelemetry.io/otel v1.26.0 // indirect

go.sum

@@ -1066,6 +1066,7 @@ github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v0.0.0-20180303142811-b89eecf5ca5d/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=

internal/gcsx/bucket_manager.go

@@ -222,6 +222,9 @@ func (bm *bucketManager) SetUpBucket(
 		bm.config.TmpObjectPrefix,
 		b)
 
+	// Fetch bucket type from storage layout api and set bucket type.
+	b.BucketType()
+
 	// Check whether this bucket works, giving the user a warning early if there
 	// is some problem.
 	{

internal/storage/bucket_handle.go

@@ -27,24 +27,59 @@ import (
 	"net/http"
 
 	"cloud.google.com/go/storage"
+	control "cloud.google.com/go/storage/control/apiv2"
+	"cloud.google.com/go/storage/control/apiv2/controlpb"
+	"github.com/googleapis/gax-go/v2"
+	"github.com/googlecloudplatform/gcsfuse/v2/internal/logger"
 	"github.com/googlecloudplatform/gcsfuse/v2/internal/storage/gcs"
 	"github.com/googlecloudplatform/gcsfuse/v2/internal/storage/storageutil"
+
 	"google.golang.org/api/googleapi"
 	"google.golang.org/api/iterator"
 )
 
 type bucketHandle struct {
 	gcs.Bucket
-	bucket     *storage.BucketHandle
-	bucketName string
-	bucketType gcs.BucketType
+	bucket        *storage.BucketHandle
+	bucketName    string
+	bucketType    gcs.BucketType
+	controlClient StorageControlClientInterface
 }
 
 func (bh *bucketHandle) Name() string {
 	return bh.bucketName
 }
 
 func (bh *bucketHandle) BucketType() gcs.BucketType {
+	var nilControlClient *control.StorageControlClient = nil
+	// Note: The first invocation of this method will be slower due to a required Google Cloud Storage (GCS) fetch.
+	// Subsequent calls will be significantly faster as the results are cached in memory.
+	// While this operation is thread-safe, parallel calls during the initial fetch can result in redundant GCS requests.
+	// To avoid this, it's advisable to call this initially while mounting.
+	if bh.bucketType == gcs.Nil {
+		if bh.controlClient == nilControlClient {
+			bh.bucketType = gcs.NonHierarchical
+			return bh.bucketType
+		}
+
+		storageLayout, err := bh.getStorageLayout()
+
+		// In case bucket does not exist, set type unknown instead of panic.
+		if err != nil {
+			bh.bucketType = gcs.Unknown
+			logger.Errorf("Error returned from GetStorageLayout: %v", err)
+			return bh.bucketType
+		}
+
+		hierarchicalNamespace := storageLayout.GetHierarchicalNamespace()
+		if hierarchicalNamespace != nil && hierarchicalNamespace.Enabled {
+			bh.bucketType = gcs.Hierarchical
+			return bh.bucketType
+		}
+
+		bh.bucketType = gcs.NonHierarchical
+	}
+
 	return bh.bucketType
 }
 
@@ -430,6 +465,19 @@ func (b *bucketHandle) ComposeObjects(ctx context.Context, req *gcs.ComposeObjec
 	return
 }
 
+// TODO: Consider adding this method to the bucket interface if additional
+// layout options are needed in the future.
+func (b *bucketHandle) getStorageLayout() (*controlpb.StorageLayout, error) {
+	var callOptions []gax.CallOption
+	stoargeLayout, err := b.controlClient.GetStorageLayout(context.Background(), &controlpb.GetStorageLayoutRequest{
+		Name:      "projects/_/buckets/" + b.bucketName + "/storageLayout",
+		Prefix:    "",
+		RequestId: "",
+	}, callOptions...)
+
+	return stoargeLayout, err
+}
+
 func isStorageConditionsNotEmpty(conditions storage.Conditions) bool {
 	return conditions != (storage.Conditions{})
 }

internal/storage/bucket_handle_test.go

@@ -22,8 +22,11 @@ import (
 	"time"
 
 	"cloud.google.com/go/storage"
+	control "cloud.google.com/go/storage/control/apiv2"
+	"cloud.google.com/go/storage/control/apiv2/controlpb"
 	"github.com/googlecloudplatform/gcsfuse/v2/internal/storage/gcs"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 )
 
@@ -1187,3 +1190,62 @@ func (testSuite *BucketHandleTest) TestComposeObjectMethodWithOneSrcObjectIsDstO
 	assert.NotNil(testSuite.T(), composedObj)
 	assert.Equal(testSuite.T(), len(ContentInTestObject)+len(ContentInTestSubObject), int(composedObj.Size))
 }
+
+func (testSuite *BucketHandleTest) TestBucketTypeForHierarchicalNameSpaceTrue() {
+	mockClient := new(MockStorageControlClient)
+	mockClient.On("GetStorageLayout", mock.Anything, mock.Anything, mock.Anything).
+		Return(&controlpb.StorageLayout{
+			HierarchicalNamespace: &controlpb.StorageLayout_HierarchicalNamespace{Enabled: true},
+		}, nil)
+	testSuite.bucketHandle.controlClient = mockClient
+
+	testSuite.bucketHandle.BucketType()
+
+	assert.Equal(testSuite.T(), gcs.Hierarchical, testSuite.bucketHandle.bucketType, "Expected Hierarchical bucket type")
+}
+
+func (testSuite *BucketHandleTest) TestBucketTypeForHierarchicalNameSpaceFalse() {
+	mockClient := new(MockStorageControlClient)
+	mockClient.On("GetStorageLayout", mock.Anything, mock.Anything, mock.Anything).
+		Return(&controlpb.StorageLayout{
+			HierarchicalNamespace: &controlpb.StorageLayout_HierarchicalNamespace{Enabled: false},
+		}, nil)
+	testSuite.bucketHandle.controlClient = mockClient
+
+	testSuite.bucketHandle.BucketType()
+
+	assert.Equal(testSuite.T(), gcs.NonHierarchical, testSuite.bucketHandle.bucketType, "Expected NonHierarchical bucket type")
+}
+
+func (testSuite *BucketHandleTest) TestBucketTypeWithError() {
+	var x *controlpb.StorageLayout
+	mockClient := new(MockStorageControlClient)
+	// Test when the client returns an error.
+	mockClient.On("GetStorageLayout", mock.Anything, mock.Anything, mock.Anything).
+		Return(x, errors.New("mocked error"))
+	testSuite.bucketHandle.controlClient = mockClient
+
+	testSuite.bucketHandle.BucketType()
+
+	assert.Equal(testSuite.T(), gcs.Unknown, testSuite.bucketHandle.bucketType, "Expected Unknown when there's an error")
+}
+
+func (testSuite *BucketHandleTest) TestBucketTypeWithHierarchicalNamespaceIsNil() {
+	mockClient := new(MockStorageControlClient)
+	mockClient.On("GetStorageLayout", mock.Anything, mock.Anything, mock.Anything).
+		Return(&controlpb.StorageLayout{}, nil)
+	testSuite.bucketHandle.controlClient = mockClient
+
+	testSuite.bucketHandle.BucketType()
+
+	assert.Equal(testSuite.T(), gcs.NonHierarchical, testSuite.bucketHandle.bucketType, "Expected NonHierarchical bucket type")
+}
+
+func (testSuite *BucketHandleTest) TestDefaultBucketTypeWithControlClientNil() {
+	var nilControlClient *control.StorageControlClient = nil
+	testSuite.bucketHandle.controlClient = nilControlClient
+
+	testSuite.bucketHandle.BucketType()
+
+	assert.Equal(testSuite.T(), gcs.NonHierarchical, testSuite.bucketHandle.bucketType, "Expected Hierarchical bucket type")
+}

internal/storage/control_client_wrapper.go

@@ -0,0 +1,38 @@
+// Copyright 2024 Google Inc. All Rights Reserved.
+//
+// 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 storage
+
+import (
+	"context"
+
+	control "cloud.google.com/go/storage/control/apiv2"
+	"cloud.google.com/go/storage/control/apiv2/controlpb"
+	"github.com/googleapis/gax-go/v2"
+)
+
+// Define the interface
+type StorageControlClientInterface interface {
+	GetStorageLayout(ctx context.Context,
+		req *controlpb.GetStorageLayoutRequest,
+		opts ...gax.CallOption) (*controlpb.StorageLayout, error)
+}
+
+type StorageControlClientWrapper struct {
+	*control.StorageControlClient
+}
+
+func (w *StorageControlClientWrapper) GetStorageLayout(ctx context.Context, req *controlpb.GetStorageLayoutRequest, opts ...gax.CallOption) (*controlpb.StorageLayout, error) {
+	return w.StorageControlClient.GetStorageLayout(ctx, req, opts...)
+}

internal/storage/fake/bucket.go

@@ -440,6 +440,10 @@ func (b *bucket) BucketType() gcs.BucketType {
 	return b.bucketType
 }
 
+func (b *bucket) FetchAndSetBucketType() {
+	b.bucketType = gcs.Unknown
+}
+
 // LOCKS_EXCLUDED(b.mu)
 func (b *bucket) ListObjects(
 	ctx context.Context,

internal/storage/gcs/bucket.go

@@ -26,8 +26,10 @@ type BucketType int
 
 // BucketType enum values.
 const (
-	NonHierarchical BucketType = iota
+	Nil BucketType = iota
+	NonHierarchical
 	Hierarchical
+	Unknown
 )
 
 // Bucket represents a GCS bucket, pre-bound with a bucket name and necessary

internal/storage/mock_control_client.go

@@ -0,0 +1,36 @@
+// Copyright 2024 Google Inc. All Rights Reserved.
+//
+// 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 storage
+
+import (
+	"context"
+
+	"cloud.google.com/go/storage/control/apiv2/controlpb"
+	"github.com/googleapis/gax-go/v2"
+	"github.com/stretchr/testify/mock"
+)
+
+// MockStorageControlClient creates a mock version of the StorageControlClient.
+type MockStorageControlClient struct {
+	mock.Mock
+}
+
+// Implement the GetStorageLayout method for the mock.
+func (m *MockStorageControlClient) GetStorageLayout(ctx context.Context,
+	req *controlpb.GetStorageLayoutRequest,
+	opts ...gax.CallOption) (*controlpb.StorageLayout, error) {
+	args := m.Called(ctx, req, opts)
+	return args.Get(0).(*controlpb.StorageLayout), args.Error(1)
+}

internal/storage/storage_handle.go

@@ -24,7 +24,6 @@ import (
 	"github.com/googleapis/gax-go/v2"
 	"github.com/googlecloudplatform/gcsfuse/v2/internal/logger"
 	mountpkg "github.com/googlecloudplatform/gcsfuse/v2/internal/mount"
-	"github.com/googlecloudplatform/gcsfuse/v2/internal/storage/gcs"
 	"github.com/googlecloudplatform/gcsfuse/v2/internal/storage/storageutil"
 	"golang.org/x/net/context"
 	option "google.golang.org/api/option"
@@ -206,10 +205,10 @@ func (sh *storageClient) BucketHandle(bucketName string, billingProject string)
 		storageBucketHandle = storageBucketHandle.UserProject(billingProject)
 	}
 
-	// TODO: Implement a method that determines and returns the appropriate bucket type
-	// when the enableHNS flag is set to true.
-	bucketType := gcs.NonHierarchical
-
-	bh = &bucketHandle{bucket: storageBucketHandle, bucketType: bucketType, bucketName: bucketName}
+	bh = &bucketHandle{
+		bucket:        storageBucketHandle,
+		bucketName:    bucketName,
+		controlClient: sh.storageControlClient,
+	}
 	return
 }

internal/storage/storage_handle_test.go

@@ -55,7 +55,7 @@ func (testSuite *StorageHandleTest) TestBucketHandleWhenBucketExistsWithEmptyBil
 
 	assert.NotNil(testSuite.T(), bucketHandle)
 	assert.Equal(testSuite.T(), TestBucketName, bucketHandle.bucketName)
-	assert.Equal(testSuite.T(), gcs.NonHierarchical, bucketHandle.bucketType)
+	assert.Equal(testSuite.T(), gcs.Nil, bucketHandle.bucketType)
 }
 
 func (testSuite *StorageHandleTest) TestBucketHandleWhenBucketDoesNotExistWithEmptyBillingProject() {
@@ -71,7 +71,7 @@ func (testSuite *StorageHandleTest) TestBucketHandleWhenBucketExistsWithNonEmpty
 
 	assert.NotNil(testSuite.T(), bucketHandle)
 	assert.Equal(testSuite.T(), TestBucketName, bucketHandle.bucketName)
-	assert.Equal(testSuite.T(), gcs.NonHierarchical, bucketHandle.bucketType)
+	assert.Equal(testSuite.T(), gcs.Nil, bucketHandle.bucketType)
 }
 
 func (testSuite *StorageHandleTest) TestBucketHandleWhenBucketDoesNotExistWithNonEmptyBillingProject() {