Kevin/add drain sub command for serve services (#16502)

* cmd/tailscale/cli: add drain subCommand for serve

This commit adds the drain subcommand for serving services. After we merge advertise and serve service as one step,
we now need a way to unadvertise service and this is it.

Updates tailscale/corp#22954

Signed-off-by: KevinLiang10 <37811973+KevinLiang10@users.noreply.github.com>

* move runServeDrain and some update regarding pr comments

Signed-off-by: KevinLiang10 <37811973+KevinLiang10@users.noreply.github.com>

* some code structure change

Signed-off-by: KevinLiang10 <37811973+KevinLiang10@users.noreply.github.com>

---------

Signed-off-by: KevinLiang10 <37811973+KevinLiang10@users.noreply.github.com>

cmd/tailscale/cli/serve_v2.go

@@ -203,6 +203,16 @@ func newServeV2Command(e *serveEnv, subcmd serveMode) *ffcli.Command {
 				Exec:       e.runServeReset,
 				FlagSet:    e.newFlags("serve-reset", nil),
 			},
+			{
+				Name:       "drain",
+				ShortUsage: fmt.Sprintf("tailscale %s drain <service>", info.Name),
+				ShortHelp:  "Drain a service from the current node",
+				LongHelp: "Make the current node no longer accept new connections for the specified service.\n" +
+					"Existing connections will continue to work until they are closed, but no new connections will be accepted.\n" +
+					"Use this command to gracefully remove a service from the current node without disrupting existing connections.\n" +
+					"<service> should be a service name (e.g., svc:my-service).",
+				Exec: e.runServeDrain,
+			},
 		},
 	}
 }
@@ -443,6 +453,44 @@ func (e *serveEnv) addServiceToPrefs(ctx context.Context, serviceName string) er
 	return err
 }
 
+func (e *serveEnv) removeServiceFromPrefs(ctx context.Context, serviceName tailcfg.ServiceName) error {
+	prefs, err := e.lc.GetPrefs(ctx)
+	if err != nil {
+		return fmt.Errorf("error getting prefs: %w", err)
+	}
+	if len(prefs.AdvertiseServices) == 0 {
+		return nil // nothing to remove
+	}
+	initialLen := len(prefs.AdvertiseServices)
+	prefs.AdvertiseServices = slices.DeleteFunc(prefs.AdvertiseServices, func(s string) bool { return s == serviceName.String() })
+	if initialLen == len(prefs.AdvertiseServices) {
+		return nil // serviceName not advertised
+	}
+	_, err = e.lc.EditPrefs(ctx, &ipn.MaskedPrefs{
+		AdvertiseServicesSet: true,
+		Prefs: ipn.Prefs{
+			AdvertiseServices: prefs.AdvertiseServices,
+		},
+	})
+	return err
+}
+
+func (e *serveEnv) runServeDrain(ctx context.Context, args []string) error {
+	if len(args) == 0 {
+		return errHelp
+	}
+	if len(args) != 1 {
+		fmt.Fprintf(Stderr, "error: invalid number of arguments\n\n")
+		return errHelp
+	}
+	svc := args[0]
+	svcName := tailcfg.ServiceName(svc)
+	if err := svcName.Validate(); err != nil {
+		return fmt.Errorf("invalid service name: %s", err)
+	}
+	return e.removeServiceFromPrefs(ctx, svcName)
+}
+
 const backgroundExistsMsg = "background configuration already exists, use `tailscale %s --%s=%d off` to remove the existing configuration"
 
 // validateConfig checks if the serve config is valid to serve the type wanted on the port.

cmd/tailscale/cli/serve_v2_test.go

@@ -1212,6 +1212,54 @@ func TestAddServiceToPrefs(t *testing.T) {
 
 }
 
+func TestRemoveServiceFromPrefs(t *testing.T) {
+	tests := []struct {
+		name          string
+		svcName       tailcfg.ServiceName
+		startServices []string
+		expected      []string
+	}{
+		{
+			name:     "remove service from empty prefs",
+			svcName:  "svc:foo",
+			expected: []string{},
+		},
+		{
+			name:          "remove existing service from prefs",
+			svcName:       "svc:foo",
+			startServices: []string{"svc:foo"},
+			expected:      []string{},
+		},
+		{
+			name:          "remove service not in prefs",
+			svcName:       "svc:bar",
+			startServices: []string{"svc:foo"},
+			expected:      []string{"svc:foo"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			lc := &fakeLocalServeClient{}
+			ctx := t.Context()
+			lc.EditPrefs(ctx, &ipn.MaskedPrefs{
+				AdvertiseServicesSet: true,
+				Prefs: ipn.Prefs{
+					AdvertiseServices: tt.startServices,
+				},
+			})
+			e := &serveEnv{lc: lc, bg: bgBoolFlag{true, false}}
+			err := e.removeServiceFromPrefs(ctx, tt.svcName)
+			if err != nil {
+				t.Fatalf("removeServiceFromPrefs(%q) returned unexpected error: %v", tt.svcName, err)
+			}
+			if !slices.Equal(lc.prefs.AdvertiseServices, tt.expected) {
+				t.Errorf("removeServiceFromPrefs(%q) = %v, want %v", tt.svcName, lc.prefs.AdvertiseServices, tt.expected)
+			}
+		})
+	}
+}
+
 func TestMessageForPort(t *testing.T) {
 	svcIPMap := tailcfg.ServiceIPMappings{
 		"svc:foo": []netip.Addr{