per-rpc quota distribution between projects with Google cloud pubsub go clients

main.go

package main

/*
Sample that overrides quota project at a _per rpc_ leve.

golang allows you to set the quota project manually using the
   https://pkg.go.dev/google.golang.org/api/option#WithQuotaProject

   flag but that flag applies to the whole client

the following snippet distributes quota between two projects

first, make sureyou don't have a default override already set

eg, incase you ever ran

gcloud auth application-default set-quota-project projectB


if you did, you would see `quota_project_id` here in this file

```
~/.config/gcloud$ more application_default_credentials.json
{
  "client_id": "764086051850-6qr4p6gpi6hn506pt8ejuq83di341hur.apps.googleusercontent.com",
  "client_secret": "...",
  "quota_project_id": "projectB",
  "refresh_token": "1//0d6redacted",
  "type": "authorized_user"
}
```

remove that key/value and safe


then you need two projects projectA,projectB  assuming you have the
 user has serviceusage.services.use permission (in roles/serviceusage.serviceUsageConsumer) on projectA and projectB

 then set the values as below


 you can use this technique to evenly distribute quota between two projects per rpc.

 admittedly, this is a rare, rare usecase (you normally just do it with the WithQuotaProject flag)
*/

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	"cloud.google.com/go/pubsub"

	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

var (
	projectID = "projectID"
	// user has serviceusage.services.use permission on mineral-minutia-820 but notfabled-ray-104117
	quotaProjects = []string{"projectA", "projectB"}
)

func GRPCQuotaInterceptor() grpc.UnaryClientInterceptor {
	return grpc.UnaryClientInterceptor(grpcUnaryQuotaInterceptor)
}

func grpcUnaryQuotaInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
	send, _ := metadata.FromOutgoingContext(ctx)
	var targetQuotaProject string
	if ctx.Value(clientMetadataKey("quotaproject")) != nil {
		if n, ok := ctx.Value(clientMetadataKey("quotaproject")).(string); ok {
			targetQuotaProject = n
		}
	} else {
		targetQuotaProject = quotaProjects[rand.Intn(len(quotaProjects))]
	}
	send.Set("X-Goog-User-Project", targetQuotaProject)
	ctx = metadata.NewOutgoingContext(ctx, send)
	err := invoker(ctx, method, req, reply, cc, opts...)

	return err
}

type clientMetadataKey string

func main() {

	rand.Seed(time.Now().UnixNano())
	ctx := context.Background()
	client, err := pubsub.NewClient(ctx, projectID, option.WithGRPCDialOption(grpc.WithUnaryInterceptor(GRPCQuotaInterceptor())))
	if err != nil {
		fmt.Printf("Could not create pubsub Client: %v", err)
		return
	}

	// this will directly apply quota to a given project
	newCtx := context.WithValue(ctx, clientMetadataKey("quotaproject"), "projectA")
	topics := client.Topics(newCtx)
	for {
		topic, err := topics.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			fmt.Printf("Error listing topics %v", err)
			return
		}
		fmt.Println(topic)
	}

	// this will distribute quota randomly
	topics = client.Topics(ctx)
	for {
		topic, err := topics.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			fmt.Printf("Error listing topics %v", err)
			return
		}
		fmt.Println(topic)
	}

}