1 /*
2  * Copyright (C) 2024 The Android Open Source Project
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 com.android.federatedcompute.services.common;
18 
19 import com.android.internal.util.Preconditions;
20 import com.android.odp.module.common.Clock;
21 import com.android.odp.module.common.MonotonicClock;
22 
23 import java.util.concurrent.atomic.AtomicBoolean;
24 import java.util.concurrent.atomic.AtomicLong;
25 
26 /** Collection of network usage statistics. */
27 public class NetworkStats {
28     private static final String TAG = NetworkStats.class.getSimpleName();
29     private final AtomicLong mTotalBytesDownloaded;
30     private final AtomicLong mTotalBytesUploaded;
31     private final AtomicBoolean mStartTimeIsSet = new AtomicBoolean(false);
32     private final AtomicBoolean mEndTimeIsSet = new AtomicBoolean(false);
33 
34     private long mDataTransferStartTime;
35 
36     private long mDataTransferDurationInMillis = 0L;
37     private final Clock mClock;
38 
NetworkStats()39     public NetworkStats() {
40         this.mTotalBytesUploaded = new AtomicLong(0);
41         this.mTotalBytesDownloaded = new AtomicLong(0);
42         this.mClock = MonotonicClock.getInstance();
43     }
44 
45     /** Adds download bytes to existing counter. */
addBytesDownloaded(long bytesDownloaded)46     public void addBytesDownloaded(long bytesDownloaded) {
47         mTotalBytesDownloaded.addAndGet(bytesDownloaded);
48     }
49 
50     /** Adds upload bytes to existing counter. */
addBytesUploaded(long bytesUploaded)51     public void addBytesUploaded(long bytesUploaded) {
52         mTotalBytesUploaded.addAndGet(bytesUploaded);
53     }
54 
55     /** Records the time of starting data transfer e.g. download/upload models. */
recordStartTimeNow()56     public void recordStartTimeNow() {
57         Preconditions.checkArgument(
58                 !mStartTimeIsSet.get(),
59                 "NetworkStats data transfer start time should only set once.");
60         mDataTransferStartTime = mClock.elapsedRealtime();
61         mStartTimeIsSet.getAndSet(true);
62     }
63 
64     /** Records the time of ending data transfer e.g. download/upload models. */
recordEndTimeNow()65     public void recordEndTimeNow() {
66         Preconditions.checkArgument(
67                 mStartTimeIsSet.get(),
68                 "NetworkStats data transfer start time should be set when record end time");
69         Preconditions.checkArgument(
70                 !mEndTimeIsSet.get(), "NetworkStats data transfer end time should only set once.");
71         mDataTransferDurationInMillis = mClock.elapsedRealtime() - mDataTransferStartTime;
72         mEndTimeIsSet.getAndSet(true);
73     }
74 
getDataTransferDurationInMillis()75     public long getDataTransferDurationInMillis() {
76         return mDataTransferDurationInMillis;
77     }
78 
79     /** Gets total download bytes. */
getTotalBytesDownloaded()80     public long getTotalBytesDownloaded() {
81         return mTotalBytesDownloaded.get();
82     }
83 
84     /** Gets total upload bytes. */
getTotalBytesUploaded()85     public long getTotalBytesUploaded() {
86         return mTotalBytesUploaded.get();
87     }
88 }
89