LogDog: Update to archival V2.

server/cmd/logdog_collector/main.go

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package main
 
 import (
         "fmt"
+        "io"
         "time"
 
         "github.com/luci/luci-go/common/auth"
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/errors"
         gcps "github.com/luci/luci-go/common/gcloud/pubsub"
         log "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/parallel"
         "github.com/luci/luci-go/server/internal/logdog/collector"
         "github.com/luci/luci-go/server/internal/logdog/collector/coordinator"
@@ skipping 71 matching lines @@
         log.Fields{
                 "subscription": sub,
         }.Infof(c, "Successfully validated Pub/Sub subscription.")
 
         st, err := a.IntermediateStorage(c)
         if err != nil {
                 return err
         }
         defer st.Close()
 
-        // Application shutdown will now operate by cancelling the Collector's
-        // shutdown Context.
-        shutdownCtx, shutdownFunc := context.WithCancel(c)
-        a.SetShutdownFunc(shutdownFunc)
-
         // Initialize our Collector service object using a caching Coordinator
         // interface.
         coord := coordinator.NewCoordinator(a.Coordinator())
         coord = coordinator.NewCache(coord, int(ccfg.StateCacheSize), ccfg.StateCacheExpiration.Duration())
 
         coll := collector.Collector{
                 Coordinator:       coord,
                 Storage:           st,
                 MaxMessageWorkers: int(ccfg.MaxMessageWorkers),
         }
         defer coll.Close()
 
         // Execute our main subscription pull loop. It will run until the supplied
         // Context is cancelled.
         psIterator, err := psSub.Pull(c)
         if err != nil {
                 log.WithError(err).Errorf(c, "Failed to create Pub/Sub iterator.")
                 return err
         }
-        defer func() {
-                log.Debugf(c, "Waiting for Pub/Sub subscription iterator to stop...")
-                psIterator.Stop()
-                log.Debugf(c, "Pub/Sub subscription iterator has stopped.")
-        }()
+        defer psIterator.Stop()
+
+        // Application shutdown will now operate by cancelling the Collector's
+        // shutdown Context.
+        a.SetShutdownFunc(psIterator.Stop)
 
         parallel.Ignore(parallel.Run(int(ccfg.MaxConcurrentMessages), func(taskC chan<- func() error) {
                 // Loop until shut down.
-                for shutdownCtx.Err() == nil {
+                for {
                         msg, err := psIterator.Next()
-                        if err != nil {
+                        switch err {
+                        case nil:
+                                taskC <- func() error {
+                                        c := log.SetField(c, "messageID", msg.ID)
+                                        msg.Done(a.processMessage(c, &coll, msg))
+                                        return nil
+                                }
+
+                        case io.EOF, context.Canceled, context.DeadlineExceeded:
+                                return
+
+                        default:
                                 log.Fields{
                                         log.ErrorKey: err,
                                         "delay":      pubsubPullErrorDelay,
                                 }.Errorf(c, "Failed to fetch Pub/Sub message, retry after delay...")
                                 clock.Sleep(c, pubsubPullErrorDelay)
-                                continue
-                        }
-
-                        taskC <- func() error {
-                                c := log.SetField(c, "messageID", msg.ID)
-                                msg.Done(a.processMessage(c, &coll, msg))
-                                return nil
                         }
                 }
         }))
 
         log.Debugf(c, "Collector finished.")
         return nil
 }
 
 // processMessage returns true if the message should be ACK'd (deleted from
 // Pub/Sub) or false if the message should not be ACK'd.
@@ skipping 33 matching lines @@
                 }.Errorf(c, "Non-transient error ingesting Pub/Sub message; ACKing.")
                 return true
         }
 }
 
 // Entry point.
 func main() {
         a := application{}
         a.Run(context.Background(), a.runCollector)
 }