Milo: Internal buildbot masters support

milo/appengine/buildbot/pubsub.go

 // Copyright 2016 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package buildbot
 
 import (
         "bytes"
         "compress/gzip"
         "compress/zlib"
         "encoding/base64"
         "encoding/json"
         "time"
 
         "github.com/luci/gae/service/datastore"
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/iotools"
         log "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/server/router"
 
         "golang.org/x/net/context"
 
         "github.com/luci/luci-go/common/tsmon/field"
         "github.com/luci/luci-go/common/tsmon/metric"
         "github.com/luci/luci-go/common/tsmon/types"
 )
 
 var (
-        // subName is the name of the pubsub subscription that milo is expecting.
+        // publicSubName is the name of the pubsub subscription that milo is expecting.
         // TODO(hinoka): This should be read from luci-config.
-        subName = "projects/luci-milo/subscriptions/buildbot-public"
+        publicSubName   = "projects/luci-milo/subscriptions/buildbot-public"
+        internalSubName = "projects/luci-milo/subscriptions/buildbot-private"
 
         // Metrics
         buildCounter = metric.NewCounter(
                 "luci/milo/buildbot_pubsub/builds",
                 "The number of buildbot builds received by Milo from PubSub",
                 types.MetricMetadata{},
                 field.Bool("internal"),
                 field.String("master"),
                 field.String("builder"),
                 field.Bool("finished"),
@@ skipping 96 matching lines @@
 
         msg := pubSubSubscription{}
         defer r.Body.Close()
         dec := json.NewDecoder(r.Body)
         if err := dec.Decode(&msg); err != nil {
                 log.WithError(err).Errorf(
                         c, "Could not decode message.  %s", err)
                 h.WriteHeader(200) // This is a hard failure, we don't want PubSub to retry.
                 return
         }
-        if msg.Subscription != subName {
+        internal := true
+        switch msg.Subscription {
+        case publicSubName:
+                internal = false
+        case internalSubName:
+                // internal = true, but that's already set.
+        default:
                 log.Errorf(
-                        c, "Subscription name %s does not match %s", msg.Subscription, subName)
+                        c, "Subscription name %s does not match %s or %s",
+                        msg.Subscription, publicSubName, internalSubName)
                 h.WriteHeader(200)
                 return
         }
         bbMsg, err := msg.GetData()
         if err != nil {
                 log.WithError(err).Errorf(c, "Could not base64 decode message %s", err)
                 h.WriteHeader(200)
                 return
         }
         builds, master, err := unmarshal(c, bbMsg)
@@ skipping 23 matching lines @@
                         if existingBuild.Finished {
                                 // Never replace a completed build.
                                 buildCounter.Add(
                                         c, 1, false, build.Master, build.Buildername, false, "Rejected")
                                 log.Debugf(
                                         c, "Found build %s/%s/%d and it's finished, skipping", build.Master, build.Buildername, build.Number)
                                 continue
                         }
                         buildExists = true
                 }
-                // Also set the finished bit.
+                // Also set the finished and internal bit.
                 build.Finished = false
                 if len(build.Times) == 2 && build.Times[1] != nil {
                         build.Finished = true
                 }
+                build.Internal = internal
                 err = ds.Put(&build)
                 if err != nil {
                         switch err {
                         case errTooBig:
                                 // This will never work, we don't want PubSub to retry.
                                 log.WithError(err).Errorf(
                                         c, "Could not save build to datastore, failing permanently")
                                 h.WriteHeader(200)
                         default:
                                 // This is transient, we do want PubSub to retry.
                                 log.WithError(err).Errorf(c, "Could not save build in datastore")
                                 h.WriteHeader(500)
                         }
                         return
                 }
                 log.Debugf(
                         c, "Saved build %s/%s/%d, it is %s",
                         build.Master, build.Buildername, build.Number, build.Finished)
                 if buildExists {
                         buildCounter.Add(
                                 c, 1, false, build.Master, build.Buildername, build.Finished, "Replaced")
                 } else {
                         buildCounter.Add(
                                 c, 1, false, build.Master, build.Buildername, build.Finished, "New")
                 }
 
         }
         if master != nil {
-                // TODO(hinoka): Internal builds.
-                err = putDSMasterJSON(c, master, false)
+                err = putDSMasterJSON(c, master, internal)
                 if err != nil {
                         log.WithError(err).Errorf(
                                 c, "Could not save master in datastore %s", err)
                         // This is transient, we do want PubSub to retry.
                         h.WriteHeader(500)
                         return
                 }
         }
         h.WriteHeader(200)
 }