Add STH caching with configurable timeouts and intelligent retry

cmd/rekor-server/app/root.go

@@ -27,6 +27,7 @@ import (
 	"github.com/go-chi/chi/v5/middleware"
 	"github.com/sigstore/rekor/pkg/api"
 	"github.com/sigstore/rekor/pkg/log"
+	"github.com/sigstore/rekor/pkg/trillianclient"
 	cose "github.com/sigstore/rekor/pkg/types/cose/v0.0.1"
 	intoto001 "github.com/sigstore/rekor/pkg/types/intoto/v0.0.1"
 	intoto002 "github.com/sigstore/rekor/pkg/types/intoto/v0.0.2"
@@ -93,6 +94,9 @@ func init() {
 	rootCmd.PersistentFlags().Uint("trillian_log_server.tlog_id", 0, "Trillian tree id")
 	rootCmd.PersistentFlags().String("trillian_log_server.sharding_config", "", "path to config file for inactive shards, in JSON or YAML")
 	rootCmd.PersistentFlags().String("trillian_log_server.grpc_default_service_config", "", "JSON string used to configure gRPC clients for communicating with Trillian")
+	rootCmd.PersistentFlags().Duration("trillian_log_server.init_latest_root_timeout", trillianclient.DefaultInitLatestRootTimeout, "timeout for fetching the latest root during client initialization")
+	rootCmd.PersistentFlags().Duration("trillian_log_server.updater_wait_timeout", trillianclient.DefaultUpdaterWaitTimeout, "timeout for STH updater polling wait operations")
+	rootCmd.PersistentFlags().Bool("trillian_log_server.cache_sth", false, "enable cached STH client with background root updates (experimental)")
 
 	rootCmd.PersistentFlags().Uint("publish_frequency", 5, "how often to publish a new checkpoint, in minutes")
 

docker-compose.yml

@@ -116,6 +116,7 @@ services:
       "--attestation_storage_bucket=file:///var/run/attestations",
       "--search_index.storage_provider=mysql",
       "--search_index.mysql.dsn=test:zaphod@tcp(mysql:3306)/test",
+      # "--trillian_log_server.cache_sth=true",
       # Uncomment this for production logging
       # "--log_type=prod",
       ]

pkg/api/api.go

@@ -109,7 +109,22 @@ func NewAPI(treeID int64) (*API, error) {
 			inactiveGRPCConfigs[r.TreeID] = *r.GRPCConfig
 		}
 	}
-	tcm := trillianclient.NewClientManager(inactiveGRPCConfigs, defaultGRPCConfig)
+
+	// Inactive shards are frozen — their trees will never advance.
+	frozenTreeIDs := make(map[int64]bool)
+	for _, r := range ranges.GetInactive() {
+		frozenTreeIDs[r.TreeID] = true
+	}
+
+	// Read timeout configuration from command line flags/config
+	clientConfig := trillianclient.Config{
+		CacheSTH:              viper.GetBool("trillian_log_server.cache_sth"),
+		InitLatestRootTimeout: viper.GetDuration("trillian_log_server.init_latest_root_timeout"),
+		UpdaterWaitTimeout:    viper.GetDuration("trillian_log_server.updater_wait_timeout"),
+		FrozenTreeIDs:         frozenTreeIDs,
+	}
+
+	tcm := trillianclient.NewClientManager(inactiveGRPCConfigs, defaultGRPCConfig, clientConfig)
 
 	roots, err := ranges.CompleteInitialization(ctx, tcm)
 	if err != nil {

pkg/sharding/ranges_test.go

@@ -654,19 +654,6 @@ func TestCompleteInitialization_Scenarios(t *testing.T) {
 		SigningSchemeOrKeyPath: keyPath,
 	}
 
-	// --- Scenario 1: Multiple Backends ---
-	s1, close1 := setupMockServer(t, mockCtl)
-	defer close1()
-	addr1 := s1.Addr
-	port1, err := strconv.Atoi(addr1[strings.LastIndex(addr1, ":")+1:])
-	require.NoError(t, err)
-
-	s2, close2 := setupMockServer(t, mockCtl)
-	defer close2()
-	addr2 := s2.Addr
-	port2, err := strconv.Atoi(addr2[strings.LastIndex(addr2, ":")+1:])
-	require.NoError(t, err)
-
 	// --- Scenario 4: Connection Failure ---
 	// Find an unused port for the connection failure test
 	lisClosed, err := net.Listen("tcp", ":0")
@@ -683,27 +670,40 @@ func TestCompleteInitialization_Scenarios(t *testing.T) {
 	}{
 		{
 			name: "Scenario 1: Multiple Backends",
-			setup: func(_ *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
+			setup: func(t *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
 				// Setup two inactive shards, each pointing to a different server
 				inactive1, _ := initializeRange(context.Background(), LogRange{TreeID: 101, SigningConfig: activeSC})
 				inactive2, _ := initializeRange(context.Background(), LogRange{TreeID: 102, SigningConfig: activeSC})
 				logRanges.inactive = Ranges{inactive1, inactive2}
 
+				// Create isolated servers for this scenario
+				sA, closeA := setupMockServer(t, mockCtl)
+				t.Cleanup(closeA)
+				addrA := sA.Addr
+				portA, err := strconv.Atoi(addrA[strings.LastIndex(addrA, ":")+1:])
+				require.NoError(t, err)
+
+				sB, closeB := setupMockServer(t, mockCtl)
+				t.Cleanup(closeB)
+				addrB := sB.Addr
+				portB, err := strconv.Atoi(addrB[strings.LastIndex(addrB, ":")+1:])
+				require.NoError(t, err)
+
 				// Mock responses from each server
 				root1 := &types.LogRootV1{TreeSize: 42}
 				rootBytes1, _ := root1.MarshalBinary()
-				s1.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(&trillian.GetLatestSignedLogRootResponse{SignedLogRoot: &trillian.SignedLogRoot{LogRoot: rootBytes1}}, nil)
+				sA.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(&trillian.GetLatestSignedLogRootResponse{SignedLogRoot: &trillian.SignedLogRoot{LogRoot: rootBytes1}}, nil).MinTimes(1)
 
 				root2 := &types.LogRootV1{TreeSize: 84}
 				rootBytes2, _ := root2.MarshalBinary()
-				s2.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(&trillian.GetLatestSignedLogRootResponse{SignedLogRoot: &trillian.SignedLogRoot{LogRoot: rootBytes2}}, nil)
+				sB.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(&trillian.GetLatestSignedLogRootResponse{SignedLogRoot: &trillian.SignedLogRoot{LogRoot: rootBytes2}}, nil).MinTimes(1)
 
 				// Configure client manager to route to the correct servers
 				grpcConfigs := map[int64]trillianclient.GRPCConfig{
-					101: {Address: "localhost", Port: uint16(port1)},
-					102: {Address: "localhost", Port: uint16(port2)},
+					101: {Address: "localhost", Port: uint16(portA)},
+					102: {Address: "localhost", Port: uint16(portB)},
 				}
-				*tcm = trillianclient.NewClientManager(grpcConfigs, trillianclient.GRPCConfig{})
+				*tcm = trillianclient.NewClientManager(grpcConfigs, trillianclient.GRPCConfig{}, trillianclient.DefaultConfig())
 			},
 			expectErr: false,
 			postCondition: func(t *testing.T, logRanges *LogRanges, roots map[int64]types.LogRootV1) {
@@ -718,17 +718,24 @@ func TestCompleteInitialization_Scenarios(t *testing.T) {
 		},
 		{
 			name: "Scenario 2: Fallback to Default Backend",
-			setup: func(_ *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
+			setup: func(t *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
 				inactive, _ := initializeRange(context.Background(), LogRange{TreeID: 201, SigningConfig: activeSC})
 				logRanges.inactive = Ranges{inactive}
 
+				// Create a dedicated default backend for this scenario
+				sDef, closeDef := setupMockServer(t, mockCtl)
+				t.Cleanup(closeDef)
+				addr := sDef.Addr
+				port, err := strconv.Atoi(addr[strings.LastIndex(addr, ":")+1:])
+				require.NoError(t, err)
+
 				root := &types.LogRootV1{TreeSize: 99}
 				rootBytes, _ := root.MarshalBinary()
-				s1.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(&trillian.GetLatestSignedLogRootResponse{SignedLogRoot: &trillian.SignedLogRoot{LogRoot: rootBytes}}, nil)
+				sDef.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(&trillian.GetLatestSignedLogRootResponse{SignedLogRoot: &trillian.SignedLogRoot{LogRoot: rootBytes}}, nil).MinTimes(1)
 
 				// No specific config for tree 201, so it should use the default
-				defaultConfig := trillianclient.GRPCConfig{Address: "localhost", Port: uint16(port1)}
-				*tcm = trillianclient.NewClientManager(map[int64]trillianclient.GRPCConfig{}, defaultConfig)
+				defaultConfig := trillianclient.GRPCConfig{Address: "localhost", Port: uint16(port)}
+				*tcm = trillianclient.NewClientManager(map[int64]trillianclient.GRPCConfig{}, defaultConfig, trillianclient.DefaultConfig())
 			},
 			expectErr: false,
 			postCondition: func(t *testing.T, logRanges *LogRanges, roots map[int64]types.LogRootV1) {
@@ -742,7 +749,9 @@ func TestCompleteInitialization_Scenarios(t *testing.T) {
 			name: "Scenario 3: No Inactive Shards",
 			setup: func(_ *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
 				logRanges.inactive = Ranges{}
-				*tcm = trillianclient.NewClientManager(nil, trillianclient.GRPCConfig{Address: "localhost", Port: uint16(port1)})
+				// No inactive shards means the client manager won't be used.
+				// Provide a no-op default config to satisfy constructor.
+				*tcm = trillianclient.NewClientManager(nil, trillianclient.GRPCConfig{Address: "localhost", Port: 0}, trillianclient.DefaultConfig())
 			},
 			expectErr: false,
 			postCondition: func(t *testing.T, logRanges *LogRanges, roots map[int64]types.LogRootV1) {
@@ -760,23 +769,30 @@ func TestCompleteInitialization_Scenarios(t *testing.T) {
 				grpcConfigs := map[int64]trillianclient.GRPCConfig{
 					401: {Address: "localhost", Port: uint16(closedAddr.Port)},
 				}
-				*tcm = trillianclient.NewClientManager(grpcConfigs, trillianclient.GRPCConfig{})
+				*tcm = trillianclient.NewClientManager(grpcConfigs, trillianclient.GRPCConfig{}, trillianclient.DefaultConfig())
 			},
 			expectErr: true,
 		},
 		{
 			name: "Scenario 5: Trillian API Error",
-			setup: func(_ *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
+			setup: func(t *testing.T, logRanges *LogRanges, tcm **trillianclient.ClientManager) {
 				inactive, _ := initializeRange(context.Background(), LogRange{TreeID: 501, SigningConfig: activeSC})
 				logRanges.inactive = Ranges{inactive}
 
+				// Create a dedicated backend that returns an error
+				sErr, closeErr := setupMockServer(t, mockCtl)
+				t.Cleanup(closeErr)
+				addr := sErr.Addr
+				port, err := strconv.Atoi(addr[strings.LastIndex(addr, ":")+1:])
+				require.NoError(t, err)
+
 				// Mock an error from the Trillian server
-				s1.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(nil, status.Error(codes.NotFound, "tree not found"))
+				sErr.Log.EXPECT().GetLatestSignedLogRoot(gomock.Any(), gomock.Any()).Return(nil, status.Error(codes.NotFound, "tree not found")).MinTimes(1)
 
 				grpcConfigs := map[int64]trillianclient.GRPCConfig{
-					501: {Address: "localhost", Port: uint16(port1)},
+					501: {Address: "localhost", Port: uint16(port)},
 				}
-				*tcm = trillianclient.NewClientManager(grpcConfigs, trillianclient.GRPCConfig{})
+				*tcm = trillianclient.NewClientManager(grpcConfigs, trillianclient.GRPCConfig{}, trillianclient.DefaultConfig())
 			},
 			expectErr: true,
 		},

pkg/trillianclient/client_interface.go

@@ -0,0 +1,67 @@
+//
+// Copyright 2026 The Sigstore Authors.
+//
+// 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 trillianclient
+
+import (
+	"context"
+
+	"github.com/google/trillian"
+	"github.com/google/trillian/types"
+	"google.golang.org/grpc/codes"
+)
+
+// ClientInterface defines the public API for interacting with a Trillian log.
+// Two implementations exist:
+//   - simpleTrillianClient: stateless, per-RPC client (default)
+//   - TrillianClient: cached STH client with background root updates (experimental, opt-in via CacheSTH)
+type ClientInterface interface {
+	AddLeaf(ctx context.Context, byteValue []byte) *Response
+	GetLatest(ctx context.Context, firstSize int64) *Response
+	GetLeafAndProofByHash(ctx context.Context, hash []byte) *Response
+	GetLeafAndProofByIndex(ctx context.Context, index int64) *Response
+	GetConsistencyProof(ctx context.Context, firstSize, lastSize int64) *Response
+	GetLeavesByRange(ctx context.Context, startIndex, count int64) *Response
+	GetLeafWithoutProof(ctx context.Context, index int64) *Response
+	Close()
+}
+
+// Response includes a status code, an optional error message, and one of the results based on the API call
+type Response struct {
+	// Status is the status code of the response
+	Status codes.Code
+	// Error contains an error on request or client failure
+	Err error
+	// GetAddResult contains the response from queueing a leaf in Trillian
+	GetAddResult *trillian.QueueLeafResponse
+	// GetLeafAndProofResult contains the response for fetching an inclusion proof and leaf
+	GetLeafAndProofResult *trillian.GetEntryAndProofResponse
+	// GetLatestResult contains the response for the latest checkpoint
+	GetLatestResult *trillian.GetLatestSignedLogRootResponse
+	// GetConsistencyProofResult contains the response for a consistency proof between two log sizes
+	GetConsistencyProofResult *trillian.GetConsistencyProofResponse
+	// GetLeavesByRangeResult contains the response for fetching a leaf without an inclusion proof
+	GetLeavesByRangeResult *trillian.GetLeavesByRangeResponse
+	// getProofResult contains the response for an inclusion proof fetched by leaf hash
+	getProofResult *trillian.GetInclusionProofByHashResponse
+}
+
+func unmarshalLogRoot(logRoot []byte) (types.LogRootV1, error) {
+	var root types.LogRootV1
+	if err := root.UnmarshalBinary(logRoot); err != nil {
+		return types.LogRootV1{}, err
+	}
+	return root, nil
+}

pkg/trillianclient/doc.go

@@ -0,0 +1,41 @@
+//
+// Copyright 2026 The Sigstore Authors.
+//
+// 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 trillianclient
+
+// Package trillianclient provides Rekor wrappers around Trillian's gRPC API.
+//
+// Two client modes are supported:
+//
+//   - simpleTrillianClient (default): stateless, per-RPC behavior with no
+//     background goroutines and no cached root state.
+//
+//   - TrillianClient (enabled with --trillian_log_server.cache_sth): cached
+//     Signed Tree Head (STH) behavior with a background updater.
+//
+// In cached mode, the client keeps an atomic snapshot of the latest verified
+// root and uses waiter channels to wake only callers whose requested tree size
+// has been reached.
+//
+// Frozen trees (inactive shards) are identified through configuration and are
+// treated specially: the client initializes once, does not start an updater,
+// and fails fast when callers request sizes that cannot be reached.
+//
+// The package exposes metrics for updater health, root advancement, and waiting
+// behavior to support operational monitoring.
+//
+// This package intentionally focuses on behavior and architecture. Any concrete
+// latency or throughput expectations depend on deployment topology, Trillian
+// configuration, and workload characteristics.