1 /*
2  * Copyright (C) 2023 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 #pragma once
18 
19 #include <jni.h>
20 
21 #include <memory>
22 
23 #include "absl/status/statusor.h"
24 #include "fcp/client/simple_task_environment.h"
25 #include "fcp/jni/jni_util.h"
26 #include "fcp/protos/plan.pb.h"
27 
28 /**
29  * A wrapper around a Java class for callbacks into Java. This is required for
30  * functionality we do not provide in c++ code, such as Clearcut logging,
31  * checking for device conditions, or querying apps for examples.
32  */
33 
34 namespace fcp {
35 namespace client {
36 namespace engine {
37 namespace jni {
38 
39 using ::google::internal::federated::plan::ExampleSelector;
40 
41 // Descriptions of Java classes we call into - full class path, and method
42 // names + JNI signatures.
43 struct SimpleTaskEnvironmentImplClassDesc {
44   static constexpr fcp::jni::JavaMethodSig kTrainingConditionsSatisfied = {
45       "trainingConditionsSatisfied", "()Z"};
46   static constexpr fcp::jni::JavaMethodSig kCreateExampleIterator = {
47       "createExampleIteratorWithContext",
48       "([B[B)Lcom/android/federatedcompute/services/training/jni/"
49       "JavaExampleIterator;"};
50 };
51 
52 class SimpleTaskEnvironmentWrapperImpl : public SimpleTaskEnvironment {
53  public:
54   explicit SimpleTaskEnvironmentWrapperImpl(JavaVM *jvm,
55                                             jobject simple_task_env);
56 
57   ~SimpleTaskEnvironmentWrapperImpl() override;
58 
59   std::string GetBaseDir() override;
60 
61   std::string GetCacheDir() override;
62 
63   bool TrainingConditionsSatisfied() override;
64 
65   absl::StatusOr<std::unique_ptr<ExampleIterator>> CreateExampleIterator(
66       const google::internal::federated::plan::ExampleSelector
67           &example_selector) override;
68 
69   absl::StatusOr<std::unique_ptr<ExampleIterator>> CreateExampleIterator(
70       const ExampleSelector &example_selector,
71       const SelectorContext &selector_context) override;
72 
73  private:
74   JavaVM *jvm_;
75   jobject jthis_;
76   jmethodID training_conditions_satisfied_id_;
77   jmethodID create_example_iterator_id_;
78 };
79 
80 }  // namespace jni
81 }  // namespace engine
82 }  // namespace client
83 }  // namespace fcp
84