agent: save client num from heartbeat message

agent/server.go

@@ -19,6 +19,7 @@ import (
 	"os"
 	"os/exec"
 	"os/signal"
+	"path/filepath"
 	"syscall"
 	"time"
 
@@ -33,6 +34,7 @@ type transporterServer struct {
 
 	databaseLogFile      *os.File
 	proxyDatabaseLogfile *os.File
+	etcdClientNumPath    string
 
 	cmd *exec.Cmd
 	pid int64
@@ -58,9 +60,10 @@ func NewServer() agentpb.TransporterServer {
 	signal.Notify(notifier, syscall.SIGINT, syscall.SIGTERM)
 
 	return &transporterServer{
-		uploadSig: make(chan struct{}, 1),
-		csvReady:  make(chan struct{}),
-		notifier:  notifier,
+		etcdClientNumPath: filepath.Join(homeDir(), "etcd-client-num"),
+		uploadSig:         make(chan struct{}, 1),
+		csvReady:          make(chan struct{}),
+		notifier:          notifier,
 	}
 }
 
@@ -226,6 +229,11 @@ func (t *transporterServer) Transfer(ctx context.Context, r *agentpb.Request) (*
 			return nil, err
 		}
 
+	case agentpb.Request_Heartbeat:
+		if err := toFile(fmt.Sprintf("%d", t.req.ClientNum), t.etcdClientNumPath); err != nil {
+			return nil, err
+		}
+
 	default:
 		return nil, fmt.Errorf("Not implemented %v", r.Operation)
 	}

agent/system_metrics.go

@@ -17,7 +17,6 @@ package agent
 import (
 	"fmt"
 	"os"
-	"path/filepath"
 	"time"
 
 	"github.com/gyuho/psn"
@@ -35,12 +34,10 @@ func startMetrics(fs *flags, t *transporterServer) error {
 		return err
 	}
 
-	epath := filepath.Join(homeDir(), "etcd-client-num")
-	if err := toFile(fmt.Sprintf("%d", t.req.ClientNum), epath); err != nil {
+	if err := toFile(fmt.Sprintf("%d", t.req.ClientNum), t.etcdClientNumPath); err != nil {
 		return err
 	}
-
-	c := psn.NewCSV(fs.systemMetricsLog, t.pid, fs.diskDevice, fs.networkInterface, epath)
+	c := psn.NewCSV(fs.systemMetricsLog, t.pid, fs.diskDevice, fs.networkInterface, t.etcdClientNumPath)
 	if err := c.Add(); err != nil {
 		return err
 	}