github.com/GoogleCloudPlatform/testgrid@v0.0.174/util/queue/persist.go (about)

     1  /*
     2  Copyright 2022 The TestGrid Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package queue
    18  
    19  import (
    20  	"context"
    21  	"encoding/json"
    22  	"errors"
    23  	"fmt"
    24  	"time"
    25  
    26  	"cloud.google.com/go/storage"
    27  	"github.com/GoogleCloudPlatform/testgrid/util/gcs"
    28  	"github.com/sirupsen/logrus"
    29  )
    30  
    31  // PersistClient contains interfaces for reading from and writing to a Path.
    32  type PersistClient interface {
    33  	gcs.Uploader
    34  	gcs.Opener
    35  }
    36  
    37  // Fixer will adjust the queue until the context expires.
    38  type Fixer func(context.Context, *Queue) error
    39  
    40  // FixPersistent persists the queue to the remote path every tick.
    41  //
    42  // The first time it will load the state. Thereafter it will save the state.
    43  // This includes restarts due to expiring contexts -- it will just load once.
    44  func FixPersistent(logr logrus.FieldLogger, client PersistClient, path gcs.Path, tick <-chan time.Time) Fixer {
    45  	var shouldSave bool
    46  	log := logr.WithField("path", path)
    47  	return func(parentCtx context.Context, q *Queue) error {
    48  		log.Debug("Using persistent state")
    49  
    50  		ctx, cancel := context.WithCancel(parentCtx)
    51  		defer cancel()
    52  		go func() { // allow a grace period for reads/writes
    53  			select {
    54  			case <-ctx.Done():
    55  				return
    56  			case <-parentCtx.Done():
    57  				timer := time.NewTimer(5 * time.Second)
    58  				select {
    59  				case <-ctx.Done():
    60  					if !timer.Stop() {
    61  						<-timer.C
    62  					}
    63  					return
    64  				case <-timer.C:
    65  					cancel()
    66  					return
    67  				}
    68  			}
    69  		}()
    70  
    71  		tryLoad := func() error {
    72  			reader, attrs, err := client.Open(ctx, path)
    73  			if errors.Is(err, storage.ErrObjectNotExist) {
    74  				log.Info("Previous persistent queue state does not exist.")
    75  				return nil
    76  			}
    77  			if err != nil {
    78  				return fmt.Errorf("open: %w", err)
    79  			}
    80  
    81  			defer reader.Close()
    82  			dec := json.NewDecoder(reader)
    83  			var whens map[string]time.Time
    84  			if err := dec.Decode(&whens); err != nil {
    85  				return fmt.Errorf("decode: %v", err)
    86  			}
    87  
    88  			current := q.Current()
    89  
    90  			for name := range whens {
    91  				if _, ok := current[name]; ok {
    92  					continue
    93  				}
    94  				delete(whens, name)
    95  			}
    96  
    97  			log.WithField("from", attrs.LastModified).Info("Loaded previous state, syncing queue.")
    98  			if err := q.FixAll(whens, false); err != nil {
    99  				return fmt.Errorf("fix all: %v", err)
   100  			}
   101  			return nil
   102  		}
   103  
   104  		logSave := true
   105  
   106  		trySave := func() error {
   107  			currently := q.Current()
   108  			buf, err := json.MarshalIndent(currently, "", "  ")
   109  			if err != nil {
   110  				return fmt.Errorf("marshal: %v", err)
   111  			}
   112  			attrs, err := client.Upload(ctx, path, buf, gcs.DefaultACL, gcs.NoCache)
   113  			if err == nil && logSave {
   114  				logSave = false
   115  				log.WithField("updated", attrs.Updated).Info("Wrote persistent state")
   116  			}
   117  			return err
   118  		}
   119  
   120  		for {
   121  			select {
   122  			case <-parentCtx.Done():
   123  				log.Debug("Stopped syncing persistent state")
   124  				return parentCtx.Err()
   125  			case <-tick:
   126  			}
   127  
   128  			if shouldSave {
   129  				log.Trace("Saving persistent state...")
   130  				if err := trySave(); err != nil {
   131  					log := log.WithError(err)
   132  					var chirp func(...interface{})
   133  					if errors.Is(err, context.Canceled) {
   134  						chirp = log.Debug
   135  					} else {
   136  						chirp = log.Warning
   137  					}
   138  					chirp("Failed to save persistent state.")
   139  					continue
   140  				}
   141  				log.Debug("Saved persistent state.")
   142  			} else {
   143  				log.Trace("Loading persistent state...")
   144  				if err := tryLoad(); err != nil {
   145  					log := log.WithError(err)
   146  					var chirp func(...interface{})
   147  					if errors.Is(err, context.Canceled) {
   148  						chirp = log.Debug
   149  					} else {
   150  						chirp = log.Warning
   151  					}
   152  					chirp("Failed to load persistent state.")
   153  					continue
   154  				}
   155  				shouldSave = true
   156  				log.Debug("Loaded persistent state.")
   157  			}
   158  		}
   159  	}
   160  }